From 95db81371a55111916188ab8d8820a1198a284f8 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 20 Jun 2024 11:22:48 +0200
Subject: [PATCH 01/97] Remove no longer needed utilities

---
 baybe/surrogates/base.py  |  8 --------
 baybe/surrogates/utils.py | 43 +--------------------------------------
 2 files changed, 1 insertion(+), 50 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 999bd7c2f..b00e5928f 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -22,7 +22,6 @@
     unstructure_base,
 )
 from baybe.serialization.mixin import SerialMixin
-from baybe.surrogates.utils import _prepare_inputs, _prepare_targets
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
@@ -77,9 +76,6 @@ def posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         """
         import torch
 
-        # Prepare the input
-        candidates = _prepare_inputs(candidates)
-
         # Evaluate the posterior distribution
         mean, covar = self._posterior(candidates)
 
@@ -147,10 +143,6 @@ def fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> Non
                 "Continuous search spaces are currently only supported by GPs."
             )
 
-        # Validate and prepare the training data
-        train_x = _prepare_inputs(train_x)
-        train_y = _prepare_targets(train_y)
-
         self._fit(searchspace, train_x, train_y)
 
     @abstractmethod
diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index 490ff3cfd..15368b02f 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -15,47 +15,6 @@
     from baybe.surrogates.base import Surrogate
 
 
-def _prepare_inputs(x: Tensor) -> Tensor:
-    """Validate and prepare the model input.
-
-    Args:
-        x: The "raw" model input.
-
-    Returns:
-        The prepared input.
-
-    Raises:
-        ValueError: If the model input is empty.
-    """
-    from baybe.utils.torch import DTypeFloatTorch
-
-    if len(x) == 0:
-        raise ValueError("The model input must be non-empty.")
-    return x.to(DTypeFloatTorch)
-
-
-def _prepare_targets(y: Tensor) -> Tensor:
-    """Validate and prepare the model targets.
-
-    Args:
-        y: The "raw" model targets.
-
-    Returns:
-        The prepared targets.
-
-    Raises:
-        NotImplementedError: If there is more than one target.
-    """
-    from baybe.utils.torch import DTypeFloatTorch
-
-    if y.shape[1] != 1:
-        raise NotImplementedError(
-            "The model currently supports only one target or multiple targets in "
-            "DESIRABILITY mode."
-        )
-    return y.to(DTypeFloatTorch)
-
-
 def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
     """Make a ``Surrogate`` class robustly handle constant training targets.
 
@@ -195,7 +154,7 @@ def _fit_new(
 
 
 def batchify(
-    posterior: Callable[[Surrogate, Tensor], tuple[Tensor, Tensor]]
+    posterior: Callable[[Surrogate, Tensor], tuple[Tensor, Tensor]],
 ) -> Callable[[Surrogate, Tensor], tuple[Tensor, Tensor]]:
     """Wrap ``Surrogate`` posterior functions to enable proper batching.
 

From 3a6ca58c8af841036233b7b652b331e75f2b91d6 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 20 Jun 2024 12:55:18 +0200
Subject: [PATCH 02/97] Move exp rep to comp rep transition point to surrogate

---
 baybe/acquisition/base.py                |  8 +--
 baybe/exceptions.py                      |  4 ++
 baybe/recommenders/pure/bayesian/base.py | 10 +---
 baybe/surrogates/base.py                 | 65 ++++++++++++++++++------
 4 files changed, 60 insertions(+), 27 deletions(-)

diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index 6b5679957..67539a0ca 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -39,8 +39,7 @@ def to_botorch(
         self,
         surrogate: Surrogate,
         searchspace: SearchSpace,
-        train_x: pd.DataFrame,
-        train_y: pd.DataFrame,
+        measurements: pd.DataFrame,
     ):
         """Create the botorch-ready representation of the function."""
         import botorch.acquisition as botorch_acqf_module
@@ -48,6 +47,9 @@ def to_botorch(
         acqf_cls = getattr(botorch_acqf_module, self.__class__.__name__)
         params_dict = filter_attributes(object=self, callable_=acqf_cls.__init__)
 
+        train_x = surrogate.transform_inputs(measurements)
+        train_y = surrogate.transform_targets(measurements)
+
         signature_params = signature(acqf_cls).parameters
         additional_params = {}
         if "model" in signature_params:
@@ -55,7 +57,7 @@ def to_botorch(
         if "best_f" in signature_params:
             additional_params["best_f"] = train_y.max().item()
         if "X_baseline" in signature_params:
-            additional_params["X_baseline"] = to_tensor(train_x)
+            additional_params["X_baseline"] = train_x
         if "mc_points" in signature_params:
             additional_params["mc_points"] = to_tensor(
                 self.get_integration_points(searchspace)  # type: ignore[attr-defined]
diff --git a/baybe/exceptions.py b/baybe/exceptions.py
index dd3010ace..a39a1c280 100644
--- a/baybe/exceptions.py
+++ b/baybe/exceptions.py
@@ -50,3 +50,7 @@ class DeprecationError(Exception):
 
 class UnidentifiedSubclassError(Exception):
     """A specified subclass cannot be found in the given class hierarchy."""
+
+
+class ModelNotTrainedError(Exception):
+    """A prediction/transformation is attempted before the model has been trained."""
diff --git a/baybe/recommenders/pure/bayesian/base.py b/baybe/recommenders/pure/bayesian/base.py
index 1e93c9ef9..cdc4f701a 100644
--- a/baybe/recommenders/pure/bayesian/base.py
+++ b/baybe/recommenders/pure/bayesian/base.py
@@ -14,7 +14,6 @@
 from baybe.searchspace import SearchSpace
 from baybe.surrogates import CustomONNXSurrogate, GaussianProcessSurrogate
 from baybe.surrogates.base import Surrogate
-from baybe.utils.dataframe import to_tensor
 
 
 @define
@@ -51,14 +50,9 @@ def _setup_botorch_acqf(
         measurements: pd.DataFrame,
     ) -> None:
         """Create the acquisition function for the current training data."""  # noqa: E501
-        # TODO: Transition point from dataframe to tensor needs to be refactored.
-        #   Currently, surrogate models operate with tensors, while acquisition
-        #   functions with dataframes.
-        train_x = searchspace.transform(measurements)
-        train_y = objective.transform(measurements)
-        self.surrogate_model._fit(searchspace, *to_tensor(train_x, train_y))
+        self.surrogate_model.fit(searchspace, objective, measurements)
         self._botorch_acqf = self.acquisition_function.to_botorch(
-            self.surrogate_model, searchspace, train_x, train_y
+            self.surrogate_model, searchspace, measurements
         )
 
     def recommend(  # noqa: D102
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index b00e5928f..76351e711 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -5,7 +5,8 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, ClassVar
 
-from attrs import define
+import pandas as pd
+from attrs import define, field
 from cattrs import override
 from cattrs.dispatch import (
     StructuredValue,
@@ -15,6 +16,8 @@
     UnstructureHook,
 )
 
+from baybe.exceptions import ModelNotTrainedError
+from baybe.objectives.base import Objective
 from baybe.searchspace import SearchSpace
 from baybe.serialization.core import (
     converter,
@@ -22,6 +25,7 @@
     unstructure_base,
 )
 from baybe.serialization.mixin import SerialMixin
+from baybe.utils.dataframe import to_tensor
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
@@ -55,29 +59,44 @@ class Surrogate(ABC, SerialMixin):
     """Class variable encoding whether or not the surrogate supports transfer
     learning."""
 
+    _input_transformd = field(init=False, default=None, eq=False)
+    """Callable preparing surrogate inputs for training/prediction.
+
+    Transforms a dataframe containing parameter configurations in experimental
+    representation to a tensor containing the same configurations in computational
+    representation. Only available after the surrogate has been fitted."""
+
+    _target_transform = field(init=False, default=None, eq=False)
+    """Callable preparing surrogate targets for training.
+
+    Transforms a dataframe containing target measurements in experimental
+    representation to a tensor containing the same measurements in computational
+    representation. Only available after the surrogate has been fitted."""
+
     def to_botorch(self) -> Model:
         """Create the botorch-ready representation of the model."""
         from baybe.surrogates._adapter import AdapterModel
 
         return AdapterModel(self)
 
-    def posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
-        """Evaluate the surrogate model at the given candidate points.
+    def transform_inputs(self, data: pd.DataFrame) -> Tensor:
+        """Transform an experimental parameter dataframe."""
+        if self._input_transform is None:
+            raise ModelNotTrainedError("The model must be trained first.")
+        return self._input_transform(data)
 
-        Args:
-            candidates: The candidate points, represented as a tensor of shape
-                ``(*t, q, d)``, where ``t`` denotes the "t-batch" shape, ``q``
-                denotes the "q-batch" shape, and ``d`` is the input dimension. For
-                more details about batch shapes, see: https://botorch.org/docs/batching
+    def transform_targets(self, data: pd.DataFrame) -> Tensor:
+        """Transform an experimental measurement dataframe."""
+        if self._target_transform is None:
+            raise ModelNotTrainedError("The model must be trained first.")
+        return self._target_transform(data)
 
-        Returns:
-            The posterior means and posterior covariance matrices of the t-batched
-            candidate points.
-        """
+    def posterior(self, candidates: pd.DataFrame) -> tuple[Tensor, Tensor]:
+        """Evaluate the surrogate model at the given candidate points."""
         import torch
 
         # Evaluate the posterior distribution
-        mean, covar = self._posterior(candidates)
+        mean, covar = self._posterior(self.transform_inputs(candidates))
 
         # Apply covariance transformation for marginal posterior models
         if not self.joint_posterior:
@@ -114,13 +133,18 @@ def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
             See :func:`baybe.surrogates.Surrogate.posterior`.
         """
 
-    def fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
+    def fit(
+        self,
+        searchspace: SearchSpace,
+        objective: Objective,
+        measurements: pd.DataFrame,
+    ) -> None:
         """Train the surrogate model on the provided data.
 
         Args:
             searchspace: The search space in which experiments are conducted.
-            train_x: The training data points.
-            train_y: The training data labels.
+            objective: The objective to be optimized.
+            measurements: The training data in experimental representation.
 
         Raises:
             ValueError: If the search space contains task parameters but the selected
@@ -143,6 +167,15 @@ def fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> Non
                 "Continuous search spaces are currently only supported by GPs."
             )
 
+        # Store context-specific transformations
+        self._input_transform = lambda x: to_tensor(searchspace.transform(x))
+        self._target_transform = lambda x: to_tensor(objective.transform(x))
+
+        # Transform and fit
+        train_x, train_y = (
+            self.transform_inputs(measurements),
+            self.transform_targets(measurements),
+        )
         self._fit(searchspace, train_x, train_y)
 
     @abstractmethod

From 023b5aa963bf7036e0eeade73a4ae3ab778bc267 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 21 Jun 2024 22:11:22 +0200
Subject: [PATCH 03/97] Transform to dataframe instead of tensor

---
 baybe/acquisition/base.py |  2 +-
 baybe/surrogates/base.py  | 18 +++++++++---------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index 67539a0ca..025a0ad40 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -57,7 +57,7 @@ def to_botorch(
         if "best_f" in signature_params:
             additional_params["best_f"] = train_y.max().item()
         if "X_baseline" in signature_params:
-            additional_params["X_baseline"] = train_x
+            additional_params["X_baseline"] = to_tensor(train_x)
         if "mc_points" in signature_params:
             additional_params["mc_points"] = to_tensor(
                 self.get_integration_points(searchspace)  # type: ignore[attr-defined]
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 76351e711..d710cc43f 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -59,18 +59,18 @@ class Surrogate(ABC, SerialMixin):
     """Class variable encoding whether or not the surrogate supports transfer
     learning."""
 
-    _input_transformd = field(init=False, default=None, eq=False)
+    _input_transform = field(init=False, default=None, eq=False)
     """Callable preparing surrogate inputs for training/prediction.
 
     Transforms a dataframe containing parameter configurations in experimental
-    representation to a tensor containing the same configurations in computational
+    representation to a corresponding dataframe containing their computational
     representation. Only available after the surrogate has been fitted."""
 
     _target_transform = field(init=False, default=None, eq=False)
     """Callable preparing surrogate targets for training.
 
     Transforms a dataframe containing target measurements in experimental
-    representation to a tensor containing the same measurements in computational
+    representation to a corresponding dataframe containing their computational
     representation. Only available after the surrogate has been fitted."""
 
     def to_botorch(self) -> Model:
@@ -79,13 +79,13 @@ def to_botorch(self) -> Model:
 
         return AdapterModel(self)
 
-    def transform_inputs(self, data: pd.DataFrame) -> Tensor:
+    def transform_inputs(self, data: pd.DataFrame) -> pd.DataFrame:
         """Transform an experimental parameter dataframe."""
         if self._input_transform is None:
             raise ModelNotTrainedError("The model must be trained first.")
         return self._input_transform(data)
 
-    def transform_targets(self, data: pd.DataFrame) -> Tensor:
+    def transform_targets(self, data: pd.DataFrame) -> pd.DataFrame:
         """Transform an experimental measurement dataframe."""
         if self._target_transform is None:
             raise ModelNotTrainedError("The model must be trained first.")
@@ -96,7 +96,7 @@ def posterior(self, candidates: pd.DataFrame) -> tuple[Tensor, Tensor]:
         import torch
 
         # Evaluate the posterior distribution
-        mean, covar = self._posterior(self.transform_inputs(candidates))
+        mean, covar = self._posterior(to_tensor(self.transform_inputs(candidates)))
 
         # Apply covariance transformation for marginal posterior models
         if not self.joint_posterior:
@@ -168,11 +168,11 @@ def fit(
             )
 
         # Store context-specific transformations
-        self._input_transform = lambda x: to_tensor(searchspace.transform(x))
-        self._target_transform = lambda x: to_tensor(objective.transform(x))
+        self._input_transform = lambda x: searchspace.transform(x)
+        self._target_transform = lambda x: objective.transform(x)
 
         # Transform and fit
-        train_x, train_y = (
+        train_x, train_y = to_tensor(
             self.transform_inputs(measurements),
             self.transform_targets(measurements),
         )

From 376ffb1fbac2d994ba87df8d7b2eed972e65a47e Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 25 Jun 2024 16:36:36 +0200
Subject: [PATCH 04/97] Update CHANGELOG.md

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4cd73f698..3bcfa53d9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `Parameter.is_numeric` has been replaced with `Parameter.is_numerical`
 - `DiscreteParameter.transform_rep_exp2comp` has been replaced with
   `DiscreteParameter.transform` 
+- `Surrogate` models now operate on dataframes in experimental representation instead of
+  tensors in computational representation
 
 ### Added
 - `Surrogate` base class now exposes a `to_botorch` method

From c0fb982297d8f7640bea8745a70847fec2ef3ca8 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 26 Jun 2024 11:24:04 +0200
Subject: [PATCH 05/97] Make recommenders operate on experimental
 representation

---
 baybe/recommenders/naive.py                   |  4 ++--
 baybe/recommenders/pure/base.py               | 22 +++++++++----------
 baybe/recommenders/pure/bayesian/botorch.py   | 18 +++++++++------
 .../pure/nonpredictive/clustering.py          |  3 ++-
 .../pure/nonpredictive/sampling.py            |  8 ++++---
 baybe/surrogates/base.py                      |  2 +-
 6 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/baybe/recommenders/naive.py b/baybe/recommenders/naive.py
index c9d10acb5..f7191fe20 100644
--- a/baybe/recommenders/naive.py
+++ b/baybe/recommenders/naive.py
@@ -122,7 +122,7 @@ def recommend(  # noqa: D102
         # Get discrete candidates. The metadata flags are ignored since the search space
         # is hybrid
         # TODO Slight BOILERPLATE CODE, see recommender.py, ll. 47+
-        _, candidates_comp = searchspace.discrete.get_candidates(
+        candidates_exp, _ = searchspace.discrete.get_candidates(
             allow_repeated_recommendations=True,
             allow_recommending_already_measured=True,
         )
@@ -147,7 +147,7 @@ def recommend(  # noqa: D102
         # Call the private function of the discrete recommender and get the indices
         disc_rec_idx = self.disc_recommender._recommend_discrete(
             subspace_discrete=searchspace.discrete,
-            candidates_comp=candidates_comp,
+            candidates_exp=candidates_exp,
             batch_size=batch_size,
         )
 
diff --git a/baybe/recommenders/pure/base.py b/baybe/recommenders/pure/base.py
index aad91ecc3..ea56b127c 100644
--- a/baybe/recommenders/pure/base.py
+++ b/baybe/recommenders/pure/base.py
@@ -50,7 +50,7 @@ def recommend(  # noqa: D102
     def _recommend_discrete(
         self,
         subspace_discrete: SubspaceDiscrete,
-        candidates_comp: pd.DataFrame,
+        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.Index:
         """Generate recommendations from a discrete search space.
@@ -58,7 +58,7 @@ def _recommend_discrete(
         Args:
             subspace_discrete: The discrete subspace from which to generate
                 recommendations.
-            candidates_comp: The computational representation of all discrete candidate
+            candidates_exp: The experimental representation of all discrete candidate
                 points to be considered.
             batch_size: The size of the recommendation batch.
 
@@ -67,14 +67,14 @@ def _recommend_discrete(
 
         Returns:
             The dataframe indices of the recommended points in the provided
-            computational representation.
+            experimental representation.
         """
         # If this method is not implemented by a child class, try to resort to hybrid
         # recommendation (with an empty subspace) instead.
         try:
             return self._recommend_hybrid(
                 searchspace=SearchSpace(discrete=subspace_discrete),
-                candidates_comp=candidates_comp,
+                candidates_exp=candidates_exp,
                 batch_size=batch_size,
             ).index
         except NotImplementedError as exc:
@@ -110,7 +110,7 @@ def _recommend_continuous(
         try:
             return self._recommend_hybrid(
                 searchspace=SearchSpace(continuous=subspace_continuous),
-                candidates_comp=pd.DataFrame(),
+                candidates_exp=pd.DataFrame(),
                 batch_size=batch_size,
             )
         except NotImplementedError as exc:
@@ -126,7 +126,7 @@ def _recommend_continuous(
     def _recommend_hybrid(
         self,
         searchspace: SearchSpace,
-        candidates_comp: pd.DataFrame,
+        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.DataFrame:
         """Generate recommendations from a hybrid search space.
@@ -138,7 +138,7 @@ def _recommend_hybrid(
         Args:
             searchspace: The hybrid search space from which to generate
                 recommendations.
-            candidates_comp: The computational representation of all discrete candidate
+            candidates_exp: The experimental representation of all discrete candidate
                 points to be considered.
             batch_size: The size of the recommendation batch.
 
@@ -175,7 +175,7 @@ def _recommend_with_discrete_parts(
 
         # Get discrete candidates
         # Repeated recommendations are always allowed for hybrid spaces
-        _, candidates_comp = searchspace.discrete.get_candidates(
+        candidates_exp, _ = searchspace.discrete.get_candidates(
             allow_repeated_recommendations=is_hybrid_space
             or self.allow_repeated_recommendations,
             allow_recommending_already_measured=is_hybrid_space
@@ -184,7 +184,7 @@ def _recommend_with_discrete_parts(
 
         # Check if enough candidates are left
         # TODO [15917]: This check is not perfectly correct.
-        if (not is_hybrid_space) and (len(candidates_comp) < batch_size):
+        if (not is_hybrid_space) and (len(candidates_exp) < batch_size):
             raise NotEnoughPointsLeftError(
                 f"Using the current settings, there are fewer than {batch_size} "
                 "possible data points left to recommend. This can be "
@@ -196,11 +196,11 @@ def _recommend_with_discrete_parts(
 
         # Get recommendations
         if is_hybrid_space:
-            rec = self._recommend_hybrid(searchspace, candidates_comp, batch_size)
+            rec = self._recommend_hybrid(searchspace, candidates_exp, batch_size)
             idxs = rec.index
         else:
             idxs = self._recommend_discrete(
-                searchspace.discrete, candidates_comp, batch_size
+                searchspace.discrete, candidates_exp, batch_size
             )
             rec = searchspace.discrete.exp_rep.loc[idxs, :]
 
diff --git a/baybe/recommenders/pure/bayesian/botorch.py b/baybe/recommenders/pure/bayesian/botorch.py
index 1479a2580..4587789c4 100644
--- a/baybe/recommenders/pure/bayesian/botorch.py
+++ b/baybe/recommenders/pure/bayesian/botorch.py
@@ -74,7 +74,7 @@ def _validate_percentage(  # noqa: DOC101, DOC103
     def _recommend_discrete(
         self,
         subspace_discrete: SubspaceDiscrete,
-        candidates_comp: pd.DataFrame,
+        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.Index:
         """Generate recommendations from a discrete search space.
@@ -82,7 +82,7 @@ def _recommend_discrete(
         Args:
             subspace_discrete: The discrete subspace from which to generate
                 recommendations.
-            candidates_comp: The computational representation of all discrete candidate
+            candidates_exp: The experimental representation of all discrete candidate
                 points to be considered.
             batch_size: The size of the recommendation batch.
 
@@ -92,7 +92,7 @@ def _recommend_discrete(
 
         Returns:
             The dataframe indices of the recommended points in the provided
-            computational representation.
+            experimental representation.
         """
         # For batch size > 1, this optimizer needs a MC acquisition function
         if batch_size > 1 and not self.acquisition_function.is_mc:
@@ -104,9 +104,9 @@ def _recommend_discrete(
         from botorch.optim import optimize_acqf_discrete
 
         # determine the next set of points to be tested
-        candidates_tensor = to_tensor(candidates_comp)
+        candidates_comp = self.surrogate_model.transform_inputs(candidates_exp)
         points, _ = optimize_acqf_discrete(
-            self._botorch_acqf, batch_size, candidates_tensor
+            self._botorch_acqf, batch_size, to_tensor(candidates_comp)
         )
 
         # retrieve the index of the points from the input dataframe
@@ -114,6 +114,7 @@ def _recommend_discrete(
         #   `SearchSpace._match_measurement_with_searchspace_indices` does, though using
         #   a simpler matching logic. When refactoring the SearchSpace class to
         #   handle continuous parameters, a corresponding utility could be extracted.
+        # IMPROVE: Maintain order of recommendations (currently lost during merge)
         idxs = pd.Index(
             pd.merge(
                 candidates_comp.reset_index(),
@@ -179,7 +180,7 @@ def _recommend_continuous(
     def _recommend_hybrid(
         self,
         searchspace: SearchSpace,
-        candidates_comp: pd.DataFrame,
+        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.DataFrame:
         """Recommend points using the ``optimize_acqf_mixed`` function of BoTorch.
@@ -193,7 +194,7 @@ def _recommend_hybrid(
 
         Args:
             searchspace: The search space in which the recommendations should be made.
-            candidates_comp: The computational representation of the candidates
+            candidates_exp: The experimental representation of the candidates
                 of the discrete subspace.
             batch_size: The size of the calculated batch.
 
@@ -214,6 +215,9 @@ def _recommend_hybrid(
         import torch
         from botorch.optim import optimize_acqf_mixed
 
+        # Transform discrete candidates
+        candidates_comp = self.surrogate_model.transform_inputs(candidates_exp)
+
         if len(candidates_comp) > 0:
             # Calculate the number of samples from the given percentage
             n_candidates = math.ceil(
diff --git a/baybe/recommenders/pure/nonpredictive/clustering.py b/baybe/recommenders/pure/nonpredictive/clustering.py
index 2d804bbf9..425d9202b 100644
--- a/baybe/recommenders/pure/nonpredictive/clustering.py
+++ b/baybe/recommenders/pure/nonpredictive/clustering.py
@@ -96,7 +96,7 @@ def _make_selection_custom(
     def _recommend_discrete(
         self,
         subspace_discrete: SubspaceDiscrete,
-        candidates_comp: pd.DataFrame,
+        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.Index:
         # See base class.
@@ -106,6 +106,7 @@ def _recommend_discrete(
         scaler = StandardScaler()
         scaler.fit(subspace_discrete.comp_rep)
 
+        candidates_comp = subspace_discrete.transform(candidates_exp)
         candidates_scaled = np.ascontiguousarray(scaler.transform(candidates_comp))
 
         # Set model parameters and perform fit
diff --git a/baybe/recommenders/pure/nonpredictive/sampling.py b/baybe/recommenders/pure/nonpredictive/sampling.py
index 7dcfa5733..dd8a1a73f 100644
--- a/baybe/recommenders/pure/nonpredictive/sampling.py
+++ b/baybe/recommenders/pure/nonpredictive/sampling.py
@@ -21,13 +21,13 @@ class RandomRecommender(NonPredictiveRecommender):
     def _recommend_hybrid(
         self,
         searchspace: SearchSpace,
-        candidates_comp: pd.DataFrame,
+        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.DataFrame:
         # See base class.
 
         if searchspace.type == SearchSpaceType.DISCRETE:
-            return candidates_comp.sample(batch_size)
+            return candidates_exp.sample(batch_size)
 
         cont_random = searchspace.continuous.sample_uniform(batch_size=batch_size)
         if searchspace.type == SearchSpaceType.CONTINUOUS:
@@ -56,7 +56,7 @@ class FPSRecommender(NonPredictiveRecommender):
     def _recommend_discrete(
         self,
         subspace_discrete: SubspaceDiscrete,
-        candidates_comp: pd.DataFrame,
+        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.Index:
         # See base class.
@@ -65,6 +65,8 @@ def _recommend_discrete(
         # TODO [Scaling]: scaling should be handled by search space object
         scaler = StandardScaler()
         scaler.fit(subspace_discrete.comp_rep)
+
+        candidates_comp = subspace_discrete.transform(candidates_exp)
         candidates_scaled = np.ascontiguousarray(scaler.transform(candidates_comp))
         ilocs = farthest_point_sampling(candidates_scaled, batch_size)
         return candidates_comp.index[ilocs]
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index d710cc43f..c15fd5247 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -168,7 +168,7 @@ def fit(
             )
 
         # Store context-specific transformations
-        self._input_transform = lambda x: searchspace.transform(x)
+        self._input_transform = lambda x: searchspace.transform(x, allow_missing=True)
         self._target_transform = lambda x: objective.transform(x)
 
         # Transform and fit

From 126bc77cec2e0630ce3101b13f765d3832407a89 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 28 Jun 2024 16:26:32 +0200
Subject: [PATCH 06/97] Add missing type hints to transform attributes

---
 baybe/surrogates/base.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index c15fd5247..316621911 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from typing import TYPE_CHECKING, ClassVar
 
 import pandas as pd
@@ -59,14 +60,18 @@ class Surrogate(ABC, SerialMixin):
     """Class variable encoding whether or not the surrogate supports transfer
     learning."""
 
-    _input_transform = field(init=False, default=None, eq=False)
+    _input_transform: Callable[[pd.DataFrame], pd.DataFrame] | None = field(
+        init=False, default=None, eq=False
+    )
     """Callable preparing surrogate inputs for training/prediction.
 
     Transforms a dataframe containing parameter configurations in experimental
     representation to a corresponding dataframe containing their computational
     representation. Only available after the surrogate has been fitted."""
 
-    _target_transform = field(init=False, default=None, eq=False)
+    _target_transform: Callable[[pd.DataFrame], pd.DataFrame] | None = field(
+        init=False, default=None, eq=False
+    )
     """Callable preparing surrogate targets for training.
 
     Transforms a dataframe containing target measurements in experimental

From c8c465c5e81f6801247737b3395407ec4689b3ea Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 28 Jun 2024 16:33:46 +0200
Subject: [PATCH 07/97] Add TODO for measurement validation

---
 baybe/surrogates/base.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 316621911..47b0f9945 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -157,6 +157,8 @@ def fit(
             NotImplementedError: When using a continuous search space and a non-GP
                 model.
         """
+        # TODO: consider adding a validation step for `measurements`
+
         # Check if transfer learning capabilities are needed
         if (searchspace.n_tasks > 1) and (not self.supports_transfer_learning):
             raise ValueError(

From 78e3938798fe3142c57a3969155136178cb079ed Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 28 Jun 2024 16:35:38 +0200
Subject: [PATCH 08/97] Add reference to docstring

---
 baybe/acquisition/base.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index 025a0ad40..cacbc77ef 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -41,7 +41,11 @@ def to_botorch(
         searchspace: SearchSpace,
         measurements: pd.DataFrame,
     ):
-        """Create the botorch-ready representation of the function."""
+        """Create the botorch-ready representation of the function.
+
+        The required structure of `measurements` is specified in
+        :meth:`babye.recommenders.base.RecommenderProtocol.recommend`.
+        """
         import botorch.acquisition as botorch_acqf_module
 
         acqf_cls = getattr(botorch_acqf_module, self.__class__.__name__)

From 602c6462f041dce0fe1f09211ea54ba43cc7cf45 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 8 Jul 2024 17:12:27 +0200
Subject: [PATCH 09/97] Refactor signatures using model context

---
 baybe/surrogates/base.py                  | 63 ++++------------
 baybe/surrogates/gaussian_process/core.py | 88 +++++++++++++++++------
 2 files changed, 80 insertions(+), 71 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 47b0f9945..2d973b572 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -4,7 +4,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import pandas as pd
 from attrs import define, field
@@ -30,10 +30,9 @@
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
+    from botorch.posteriors import Posterior
     from torch import Tensor
 
-# Define constants
-_MIN_VARIANCE = 1e-6
 
 _ONNX_ENCODING = "latin-1"
 """Constant signifying the encoding for onnx byte strings in pretrained models.
@@ -96,47 +95,22 @@ def transform_targets(self, data: pd.DataFrame) -> pd.DataFrame:
             raise ModelNotTrainedError("The model must be trained first.")
         return self._target_transform(data)
 
-    def posterior(self, candidates: pd.DataFrame) -> tuple[Tensor, Tensor]:
+    def posterior(self, candidates: pd.DataFrame) -> Posterior:
         """Evaluate the surrogate model at the given candidate points."""
-        import torch
-
-        # Evaluate the posterior distribution
-        mean, covar = self._posterior(to_tensor(self.transform_inputs(candidates)))
-
-        # Apply covariance transformation for marginal posterior models
-        if not self.joint_posterior:
-            # Convert to tensor containing covariance matrices
-            covar = torch.diag_embed(covar)
-
-        # Add small diagonal variances for numerical stability
-        covar.add_(torch.eye(covar.shape[-1]) * _MIN_VARIANCE)
-
-        return mean, covar
+        return self._posterior(to_tensor(self.transform_inputs(candidates)))
 
     @abstractmethod
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
-        """Perform the actual posterior evaluation logic.
-
-        In contrast to its public counterpart
-        :func:`baybe.surrogates.Surrogate.posterior`, no data
-        validation/transformation is carried out but only the raw posterior computation
-        is conducted.
-
-        Note that the public ``posterior`` method *always* returns a full covariance
-        matrix. By contrast, this method may return either a covariance matrix or a
-        tensor of marginal variances, depending on the models ``joint_posterior``
-        flag. The optional conversion to a covariance matrix is handled by the public
-        method.
+    def _posterior(self, candidates: Tensor) -> Posterior:
+        """Perform the actual posterior evaluation logic."""
 
-        See :func:`baybe.surrogates.Surrogate.posterior` for details on the
-        parameters.
+    @staticmethod
+    def _get_model_context(searchspace: SearchSpace, objective: Objective) -> Any:
+        """Get the surrogate-specific context for model fitting.
 
-        Args:
-            candidates: The candidates.
-
-        Returns:
-            See :func:`baybe.surrogates.Surrogate.posterior`.
+        By default, no context is created. If context is required, subclasses are
+        expected to override this method.
         """
+        return None
 
     def fit(
         self,
@@ -183,18 +157,11 @@ def fit(
             self.transform_inputs(measurements),
             self.transform_targets(measurements),
         )
-        self._fit(searchspace, train_x, train_y)
+        self._fit(train_x, train_y, self._get_model_context(searchspace, objective))
 
     @abstractmethod
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
-        """Perform the actual fitting logic.
-
-        In contrast to its public counterpart :func:`baybe.surrogates.Surrogate.fit`,
-        no data validation/transformation is carried out but only the raw fitting
-        operation is conducted.
-
-        See :func:`baybe.surrogates.Surrogate.fit` for details on the parameters.
-        """
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any = None) -> None:
+        """Perform the actual fitting logic."""
 
 
 def _make_hook_decode_onnx_str(
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index f22591294..cfd5b99bb 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -5,8 +5,10 @@
 from typing import TYPE_CHECKING, ClassVar
 
 from attrs import define, field
+from attrs.validators import instance_of
 
-from baybe.searchspace import SearchSpace
+from baybe.objective import Objective
+from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.base import Surrogate
 from baybe.surrogates.gaussian_process.kernel_factory import (
     KernelFactory,
@@ -23,9 +25,50 @@
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
+    from botorch.posteriors import Posterior
     from torch import Tensor
 
 
+@define
+class _ModelContext:
+    """Model context for :class:`GaussianProcessSurrogate`."""
+
+    searchspace: SearchSpace = field(validator=instance_of(SearchSpace))
+    """The search space the model is trained on."""
+
+    @property
+    def task_idx(self) -> int | None:
+        """The computational column index of the task parameter, if available."""
+        return self.searchspace.task_idx
+
+    @property
+    def is_multitask(self) -> bool:
+        """Indicates if model is to be operated in a multi-task context."""
+        return self.n_task_dimensions > 0
+
+    @property
+    def n_task_dimensions(self) -> int:
+        """The number of task dimensions."""
+        # TODO: Generalize to multiple task parameters
+        return 1 if self.task_idx is not None else 0
+
+    @property
+    def n_tasks(self) -> int:
+        """The number of tasks."""
+        return self.searchspace.n_tasks
+
+    @property
+    def parameter_bounds(self) -> Tensor:
+        """Get the search space parameter bounds in BoTorch Format."""
+        import torch
+
+        return torch.from_numpy(self.searchspace.param_bounds_comp)
+
+    def get_numerical_indices(self, n_inputs: int) -> list[int]:
+        """Get the indices of the regular numerical model inputs."""
+        return [i for i in range(n_inputs) if i != self.task_idx]
+
+
 @define
 class GaussianProcessSurrogate(Surrogate):
     """A Gaussian process surrogate model."""
@@ -65,33 +108,32 @@ def to_botorch(self) -> Model:  # noqa: D102
 
         return self._model
 
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    @staticmethod
+    def _get_model_context(
+        searchspace: SearchSpace, objective: Objective
+    ) -> _ModelContext:
         # See base class.
-        posterior = self._model.posterior(candidates)
-        return posterior.mvn.mean, posterior.mvn.covariance_matrix
+        return _ModelContext(searchspace=searchspace)
 
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
+    def _posterior(self, candidates: Tensor) -> Posterior:
+        # See base class.
+        return self._model.posterior(candidates)
+
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None:
         # See base class.
 
         import botorch
         import gpytorch
         import torch
 
-        # identify the indexes of the task and numeric dimensions
-        # TODO: generalize to multiple task parameters
-        task_idx = searchspace.task_idx
-        n_task_params = 1 if task_idx is not None else 0
-        numeric_idxs = [i for i in range(train_x.shape[1]) if i != task_idx]
-
-        # get the input bounds from the search space in BoTorch Format
-        bounds = torch.from_numpy(searchspace.param_bounds_comp)
-        # TODO: use target value bounds when explicitly provided
+        numerical_idxs = context.get_numerical_indices(train_x.shape[-1])
 
         # define the input and outcome transforms
         # TODO [Scaling]: scaling should be handled by search space object
         input_transform = botorch.models.transforms.Normalize(
-            train_x.shape[1], bounds=bounds, indices=numeric_idxs
+            train_x.shape[1], bounds=context.parameter_bounds, indices=numerical_idxs
         )
+        # TODO: use target value bounds when explicitly provided
         outcome_transform = botorch.models.transforms.Standardize(train_y.shape[1])
 
         # extract the batch shape of the training data
@@ -102,26 +144,26 @@ def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> No
 
         # define the covariance module for the numeric dimensions
         base_covar_module = self.kernel_factory(
-            searchspace, train_x, train_y
+            context.searchspace, train_x, train_y
         ).to_gpytorch(
-            ard_num_dims=train_x.shape[-1] - n_task_params,
-            active_dims=numeric_idxs,
+            ard_num_dims=train_x.shape[-1] - context.n_task_dimensions,
+            active_dims=numerical_idxs,
             batch_shape=batch_shape,
         )
 
         # create GP covariance
-        if task_idx is None:
+        if not context.is_multitask:
             covar_module = base_covar_module
         else:
             task_covar_module = gpytorch.kernels.IndexKernel(
-                num_tasks=searchspace.n_tasks,
-                active_dims=task_idx,
-                rank=searchspace.n_tasks,  # TODO: make controllable
+                num_tasks=context.n_tasks,
+                active_dims=context.task_idx,
+                rank=context.n_tasks,  # TODO: make controllable
             )
             covar_module = base_covar_module * task_covar_module
 
         # create GP likelihood
-        noise_prior = _default_noise_factory(searchspace, train_x, train_y)
+        noise_prior = _default_noise_factory(context.searchspace, train_x, train_y)
         likelihood = gpytorch.likelihoods.GaussianLikelihood(
             noise_prior=noise_prior[0].to_gpytorch(), batch_shape=batch_shape
         )

From e62d4e2aaa0548ce88441ee7fdd172de247dd5b2 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 9 Jul 2024 10:07:18 +0200
Subject: [PATCH 10/97] Add GaussianSurrogate base class

---
 baybe/surrogates/_adapter.py      |  6 +-----
 baybe/surrogates/base.py          | 30 +++++++++++++++++++++++++++++-
 baybe/surrogates/custom.py        | 14 ++++++--------
 baybe/surrogates/linear.py        |  9 ++++-----
 baybe/surrogates/naive.py         | 11 +++++------
 baybe/surrogates/ngboost.py       |  9 ++++-----
 baybe/surrogates/random_forest.py |  9 ++++-----
 7 files changed, 53 insertions(+), 35 deletions(-)

diff --git a/baybe/surrogates/_adapter.py b/baybe/surrogates/_adapter.py
index 6241ac37f..742a893b9 100644
--- a/baybe/surrogates/_adapter.py
+++ b/baybe/surrogates/_adapter.py
@@ -3,10 +3,8 @@
 from collections.abc import Callable
 from typing import Any
 
-import gpytorch.distributions
 from botorch.models.gpytorch import Model
 from botorch.posteriors import Posterior
-from botorch.posteriors.gpytorch import GPyTorchPosterior
 from torch import Tensor
 
 from baybe.surrogates.base import Surrogate
@@ -41,6 +39,4 @@ def posterior(  # noqa: D102
         **kwargs: Any,
     ) -> Posterior:
         # See base class.
-        mean, var = self._surrogate.posterior(X)
-        mvn = gpytorch.distributions.MultivariateNormal(mean, var)
-        return GPyTorchPosterior(mvn)
+        return self._surrogate._posterior(X)
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 2d973b572..01d6a3677 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -30,7 +30,7 @@
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
-    from botorch.posteriors import Posterior
+    from botorch.posteriors import GPyTorchPosterior, Posterior
     from torch import Tensor
 
 
@@ -164,6 +164,34 @@ def _fit(self, train_x: Tensor, train_y: Tensor, context: Any = None) -> None:
         """Perform the actual fitting logic."""
 
 
+@define(slots=False)
+class GaussianSurrogate(Surrogate, ABC):
+    """A surrogate model providing Gaussian posterior estimates."""
+
+    def _posterior(self, candidates: Tensor) -> GPyTorchPosterior:
+        # See base class.
+
+        import torch
+        from botorch.posteriors import GPyTorchPosterior
+        from gpytorch.distributions import MultivariateNormal
+
+        # Construct the Gaussian posterior from the estimated first and second moment
+        mean, var = self._estimate_moments(candidates)
+        if not self.joint_posterior:
+            var = torch.diag_embed(var)
+        mvn = MultivariateNormal(mean, var)
+        return GPyTorchPosterior(mvn)
+
+    @abstractmethod
+    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+        """Estimate first and second moments of the Gaussian posterior.
+
+        The second moment may either be a 1-D tensor of marginal variances for the
+        candidates or a 2-D tensor representing a full covariance matrix over all
+        candidates, depending on the ``joint_posterior`` flag of the model.
+        """
+
+
 def _make_hook_decode_onnx_str(
     raw_unstructure_hook: UnstructureHook
 ) -> UnstructureHook:
diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py
index f33233a3a..9ed19e6c8 100644
--- a/baybe/surrogates/custom.py
+++ b/baybe/surrogates/custom.py
@@ -10,7 +10,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from attrs import define, field, validators
 
@@ -24,7 +24,7 @@
 )
 from baybe.searchspace import SearchSpace
 from baybe.serialization.core import block_serialization_hook, converter
-from baybe.surrogates.base import Surrogate
+from baybe.surrogates.base import GaussianSurrogate, Surrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import validate_custom_architecture_cls
 from baybe.utils.numerical import DTypeFloatONNX
@@ -66,10 +66,8 @@ class CustomArchitectureSurrogate(Surrogate):
             def __init__(self, *args, **kwargs):
                 self._model = model_cls(*args, **kwargs)
 
-            def _fit(
-                self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-            ) -> None:
-                return self._model._fit(searchspace, train_x, train_y)
+            def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+                return self._model._fit(train_x, train_y, context)
 
             def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
                 return self._model._posterior(candidates)
@@ -113,7 +111,7 @@ def __get_attribute__(self, attr):
 
 
 @define(kw_only=True)
-class CustomONNXSurrogate(Surrogate):
+class CustomONNXSurrogate(GaussianSurrogate):
     """A wrapper class for custom pretrained surrogate models.
 
     Note that these surrogates cannot be retrained.
@@ -149,7 +147,7 @@ def default_model(self) -> ort.InferenceSession:
             raise ValueError("Invalid ONNX string") from exc
 
     @batchify
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         import torch
 
         from baybe.utils.torch import DTypeFloatTorch
diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py
index 9922c7611..463169d71 100644
--- a/baybe/surrogates/linear.py
+++ b/baybe/surrogates/linear.py
@@ -13,8 +13,7 @@
 from attr import define, field
 from sklearn.linear_model import ARDRegression
 
-from baybe.searchspace import SearchSpace
-from baybe.surrogates.base import Surrogate
+from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import autoscale, batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
 
@@ -25,7 +24,7 @@
 @catch_constant_targets
 @autoscale
 @define(slots=False)
-class BayesianLinearSurrogate(Surrogate):
+class BayesianLinearSurrogate(GaussianSurrogate):
     """A Bayesian linear regression surrogate model."""
 
     # Class variables
@@ -47,7 +46,7 @@ class BayesianLinearSurrogate(Surrogate):
     """The actual model."""
 
     @batchify
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
@@ -61,7 +60,7 @@ def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
 
         return mean, var
 
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         # See base class.
         self._model = ARDRegression(**(self.model_params))
         self._model.fit(train_x, train_y.ravel())
diff --git a/baybe/surrogates/naive.py b/baybe/surrogates/naive.py
index 2c370c1df..629ea5597 100644
--- a/baybe/surrogates/naive.py
+++ b/baybe/surrogates/naive.py
@@ -2,12 +2,11 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from attr import define, field
 
-from baybe.searchspace import SearchSpace
-from baybe.surrogates.base import Surrogate
+from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify
 
 if TYPE_CHECKING:
@@ -15,7 +14,7 @@
 
 
 @define
-class MeanPredictionSurrogate(Surrogate):
+class MeanPredictionSurrogate(GaussianSurrogate):
     """A trivial surrogate model.
 
     It provides the average value of the training targets
@@ -34,7 +33,7 @@ class MeanPredictionSurrogate(Surrogate):
     """The estimated posterior mean value of the training targets."""
 
     @batchify
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
@@ -44,6 +43,6 @@ def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         var = torch.ones(len(candidates))
         return mean, var
 
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         # See base class.
         self._model = train_y.mean().item()
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 488e3fabf..96db80a7e 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -13,8 +13,7 @@
 from attr import define, field
 from ngboost import NGBRegressor
 
-from baybe.searchspace import SearchSpace
-from baybe.surrogates.base import Surrogate
+from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import autoscale, batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
 
@@ -25,7 +24,7 @@
 @catch_constant_targets
 @autoscale
 @define(slots=False)
-class NGBoostSurrogate(Surrogate):
+class NGBoostSurrogate(GaussianSurrogate):
     """A natural-gradient-boosting surrogate model."""
 
     # Class variables
@@ -53,7 +52,7 @@ def __attrs_post_init__(self):
         self.model_params = {**self._default_model_params, **self.model_params}
 
     @batchify
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
@@ -67,6 +66,6 @@ def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
 
         return mean, var
 
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         # See base class.
         self._model = NGBRegressor(**(self.model_params)).fit(train_x, train_y.ravel())
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 78a9a5ed9..407730f14 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -14,8 +14,7 @@
 from attr import define, field
 from sklearn.ensemble import RandomForestRegressor
 
-from baybe.searchspace import SearchSpace
-from baybe.surrogates.base import Surrogate
+from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import autoscale, batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
 
@@ -26,7 +25,7 @@
 @catch_constant_targets
 @autoscale
 @define(slots=False)
-class RandomForestSurrogate(Surrogate):
+class RandomForestSurrogate(GaussianSurrogate):
     """A random forest surrogate model."""
 
     # Class variables
@@ -48,7 +47,7 @@ class RandomForestSurrogate(Surrogate):
     """The actual model."""
 
     @batchify
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
@@ -72,7 +71,7 @@ def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
 
         return mean, var
 
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         # See base class.
         self._model = RandomForestRegressor(**(self.model_params))
         self._model.fit(train_x, train_y.ravel())

From e708fd5d0cedde4e23081841492438d082f30bbe Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 9 Jul 2024 10:21:29 +0200
Subject: [PATCH 11/97] Update constant target catching

---
 baybe/surrogates/utils.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index 15368b02f..7e35ebbec 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -4,12 +4,13 @@
 
 from collections.abc import Callable
 from functools import wraps
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from baybe.scaler import DefaultScaler
 from baybe.searchspace import SearchSpace
 
 if TYPE_CHECKING:
+    from botorch.posteriors import Posterior
     from torch import Tensor
 
     from baybe.surrogates.base import Surrogate
@@ -54,7 +55,7 @@ def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
     _fit_original = cls._fit
     _posterior_original = cls._posterior
 
-    def _posterior_new(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _posterior_new(self, candidates: Tensor) -> Posterior:
         # Alternative model fallback
         if hasattr(self, injected_model_attr_name):
             return getattr(self, injected_model_attr_name)._posterior(candidates)
@@ -62,9 +63,7 @@ def _posterior_new(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         # Regular operation
         return _posterior_original(self, candidates)
 
-    def _fit_new(
-        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> None:
+    def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         if not (train_y.ndim == 2 and train_y.shape[-1] == 1):
             raise NotImplementedError(
                 "The current logic is only implemented for single-target surrogates."
@@ -73,7 +72,7 @@ def _fit_new(
         # Alternative model fallback
         if train_y.numel() == 1 or train_y.std() < std_threshold:
             model = MeanPredictionSurrogate()
-            model._fit(searchspace, train_x, train_y)
+            model._fit(train_x, train_y, context)
             try:
                 setattr(self, injected_model_attr_name, model)
             except AttributeError as ex:
@@ -86,7 +85,7 @@ def _fit_new(
         else:
             if hasattr(self, injected_model_attr_name):
                 delattr(self, injected_model_attr_name)
-            _fit_original(self, searchspace, train_x, train_y)
+            _fit_original(self, train_x, train_y, context)
 
     # Replace the methods
     cls._posterior = _posterior_new

From 9140b0705f5afe323f66c4652c5c8bb5f2b3504f Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 9 Jul 2024 10:44:12 +0200
Subject: [PATCH 12/97] Store constant target fallback models outside of
 surrogate instances

This enables reactivating slots
---
 baybe/surrogates/base.py          |  4 ++--
 baybe/surrogates/linear.py        |  2 +-
 baybe/surrogates/ngboost.py       |  2 +-
 baybe/surrogates/random_forest.py |  2 +-
 baybe/surrogates/utils.py         | 38 ++++++++++---------------------
 5 files changed, 17 insertions(+), 31 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 01d6a3677..aa7ffc582 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -47,7 +47,7 @@
 """
 
 
-@define(slots=False)
+@define
 class Surrogate(ABC, SerialMixin):
     """Abstract base class for all surrogate models."""
 
@@ -164,7 +164,7 @@ def _fit(self, train_x: Tensor, train_y: Tensor, context: Any = None) -> None:
         """Perform the actual fitting logic."""
 
 
-@define(slots=False)
+@define
 class GaussianSurrogate(Surrogate, ABC):
     """A surrogate model providing Gaussian posterior estimates."""
 
diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py
index 463169d71..c84f32ab5 100644
--- a/baybe/surrogates/linear.py
+++ b/baybe/surrogates/linear.py
@@ -23,7 +23,7 @@
 
 @catch_constant_targets
 @autoscale
-@define(slots=False)
+@define
 class BayesianLinearSurrogate(GaussianSurrogate):
     """A Bayesian linear regression surrogate model."""
 
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 96db80a7e..4a64f82db 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -23,7 +23,7 @@
 
 @catch_constant_targets
 @autoscale
-@define(slots=False)
+@define
 class NGBoostSurrogate(GaussianSurrogate):
     """A natural-gradient-boosting surrogate model."""
 
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 407730f14..e4064b985 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -24,7 +24,7 @@
 
 @catch_constant_targets
 @autoscale
-@define(slots=False)
+@define
 class RandomForestSurrogate(GaussianSurrogate):
     """A random forest surrogate model."""
 
diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index 7e35ebbec..b973baa6c 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -14,6 +14,13 @@
     from torch import Tensor
 
     from baybe.surrogates.base import Surrogate
+    from baybe.surrogates.naive import MeanPredictionSurrogate
+
+
+_constant_target_model_store: dict[int, MeanPredictionSurrogate] = {}
+"""Dictionary for storing constant target fallback models. Keys are the IDs of the
+surrogate models that temporarily have a fallback attached because they were
+trained on constant training targets. Values are the corresponding fallback models."""
 
 
 def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
@@ -25,30 +32,16 @@ def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
 
     The modified class handles the above cases separately from "regular operation"
     by resorting to a :class:`baybe.surrogates.naive.MeanPredictionSurrogate`,
-    which is stored as an additional temporary attribute in its objects.
+    which is stored outside the model in a dictionary maintained by this decorator.
 
     Args:
         cls: The :class:`baybe.surrogates.base.Surrogate` to be augmented.
         std_threshold: The standard deviation threshold below which operation is
             switched to the alternative model.
 
-    Raises:
-        ValueError: If the class already contains an attribute with the same name
-            as the temporary attribute to be added.
-
     Returns:
         The modified class.
     """
-    # Name of the attribute added to store the alternative model
-    injected_model_attr_name = "_constant_target_model"
-
-    if injected_model_attr_name in (attr.name for attr in cls.__attrs_attrs__):
-        raise ValueError(
-            f"Cannot apply '{catch_constant_targets.__name__}' because "
-            f"'{cls.__name__}' already has an attribute '{injected_model_attr_name}' "
-            f"defined."
-        )
-
     from baybe.surrogates.naive import MeanPredictionSurrogate
 
     # References to original methods
@@ -57,8 +50,8 @@ def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
 
     def _posterior_new(self, candidates: Tensor) -> Posterior:
         # Alternative model fallback
-        if hasattr(self, injected_model_attr_name):
-            return getattr(self, injected_model_attr_name)._posterior(candidates)
+        if constant_target_model := _constant_target_model_store.get(id(self), None):
+            return constant_target_model._posterior(candidates)
 
         # Regular operation
         return _posterior_original(self, candidates)
@@ -73,18 +66,11 @@ def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         if train_y.numel() == 1 or train_y.std() < std_threshold:
             model = MeanPredictionSurrogate()
             model._fit(train_x, train_y, context)
-            try:
-                setattr(self, injected_model_attr_name, model)
-            except AttributeError as ex:
-                raise TypeError(
-                    f"'{catch_constant_targets.__name__}' is only applicable to "
-                    f"non-slotted classes but '{cls.__name__}' is a slotted class."
-                ) from ex
+            _constant_target_model_store[id(self)] = model
 
         # Regular operation
         else:
-            if hasattr(self, injected_model_attr_name):
-                delattr(self, injected_model_attr_name)
+            _constant_target_model_store.pop(id(self), None)
             _fit_original(self, train_x, train_y, context)
 
     # Replace the methods

From d094a1a6c0d2c87439c08c6981233afe7f427496 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 9 Jul 2024 10:48:23 +0200
Subject: [PATCH 13/97] Catch passing of unimplemented posterior options

---
 baybe/surrogates/_adapter.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/baybe/surrogates/_adapter.py b/baybe/surrogates/_adapter.py
index 742a893b9..a766446bd 100644
--- a/baybe/surrogates/_adapter.py
+++ b/baybe/surrogates/_adapter.py
@@ -39,4 +39,12 @@ def posterior(  # noqa: D102
         **kwargs: Any,
     ) -> Posterior:
         # See base class.
+        if (
+            (output_indices is not None)
+            or observation_noise
+            or (posterior_transform is not None)
+        ):
+            raise NotImplementedError(
+                "The optional model posterior arguments are not yet implemented."
+            )
         return self._surrogate._posterior(X)

From 5f095de0b4ce6e2f86428cd68e6283101f9894ed Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 9 Jul 2024 13:12:42 +0200
Subject: [PATCH 14/97] Update CHANGELOG.md

---
 CHANGELOG.md | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3bcfa53d9..2ea783d7a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `DiscreteParameter.transform` 
 - `Surrogate` models now operate on dataframes in experimental representation instead of
   tensors in computational representation
+- `Surrogate.posterior` models now returns a `Posterior` object
 
 ### Added
 - `Surrogate` base class now exposes a `to_botorch` method
@@ -31,6 +32,7 @@ _ `_optional` subpackage for managing optional dependencies
 - `register_hooks` utility enabling user-defined augmentation of arbitrary callables
 - `transform` methods of `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous`
   now take additional `allow_missing` and `allow_extra` keyword arguments
+- `GaussianSurrogate` base class for surrogate models with Gaussian posteriors
 
 ### Changed
 - Passing an `Objective` to `Campaign` is now optional
@@ -39,10 +41,15 @@ _ `_optional` subpackage for managing optional dependencies
 - Sampling methods in `qNIPV` and `BotorchRecommender` are now specified via 
   `DiscreteSamplingMethod` enum
 - `Interval` class now supports degenerate intervals containing only one element
+- Context information required by `Surrogate` models is now cleanly encapsulated into
+  a `context` object passed to `Surrogate._fit`
+- Fallback models created by `catch_constant_targets` are stored outside of surrogate
 
 ### Removed
 - Support for Python 3.9 removed due to new [BoTorch requirements](https://github.com/pytorch/botorch/pull/2293) 
   and guidelines from [Scientific Python](https://scientific-python.org/specs/spec-0000/)
+- `register_custom_architecture` decorator
+- `Scalar` and `DefaultScaler` classes
 
 ### Fixed
 - `sequential` flag of `SequentialGreedyRecommender` is now set to `True`
@@ -57,6 +64,8 @@ _ `_optional` subpackage for managing optional dependencies
 - Passing a dataframe via the `data` argument to the `transform` methods of
   `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous` is no longer possible.
   The dataframe must now be passed as positional argument.
+- Role of `register_custom_architecture` has been taken over by
+  `baybe.surrogates.base.SurrogateProtocol`
 
 ## [0.9.1] - 2024-06-04
 ### Changed

From 32e2ef5d906cf35884d71755baeda7cc12ca93ff Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 11 Jul 2024 08:41:50 +0200
Subject: [PATCH 15/97] Add docstrings to override methods in decorator

---
 baybe/surrogates/utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index b973baa6c..b95752ac5 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -49,6 +49,7 @@ def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
     _posterior_original = cls._posterior
 
     def _posterior_new(self, candidates: Tensor) -> Posterior:
+        """Use fallback model if it exists, otherwise call original posterior."""
         # Alternative model fallback
         if constant_target_model := _constant_target_model_store.get(id(self), None):
             return constant_target_model._posterior(candidates)
@@ -57,6 +58,7 @@ def _posterior_new(self, candidates: Tensor) -> Posterior:
         return _posterior_original(self, candidates)
 
     def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+        """Original fit but with fallback model creation for constant targets."""
         if not (train_y.ndim == 2 and train_y.shape[-1] == 1):
             raise NotImplementedError(
                 "The current logic is only implemented for single-target surrogates."

From d9aefe5616c03fbfd941971ae447bf0464ea2be0 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 9 Jul 2024 10:20:06 +0200
Subject: [PATCH 16/97] Remove current scaling functionality

---
 baybe/scaler.py                   | 132 ------------------------------
 baybe/surrogates/linear.py        |   3 +-
 baybe/surrogates/ngboost.py       |   3 +-
 baybe/surrogates/random_forest.py |   3 +-
 baybe/surrogates/utils.py         |  61 --------------
 mypy.ini                          |   1 -
 6 files changed, 3 insertions(+), 200 deletions(-)
 delete mode 100644 baybe/scaler.py

diff --git a/baybe/scaler.py b/baybe/scaler.py
deleted file mode 100644
index 19fc7119c..000000000
--- a/baybe/scaler.py
+++ /dev/null
@@ -1,132 +0,0 @@
-"""Functionality for data scaling."""
-
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from collections.abc import Callable
-from typing import TYPE_CHECKING
-
-import pandas as pd
-
-from baybe.utils.dataframe import to_tensor
-
-if TYPE_CHECKING:
-    from torch import Tensor
-
-    _ScaleFun = Callable[[Tensor], Tensor]
-
-
-class Scaler(ABC):
-    """Abstract base class for all scalers.
-
-    Args:
-        searchspace: The search space that should be scaled.
-    """
-
-    type: str
-    """Class variable encoding the type of the scaler."""
-
-    SUBCLASSES: dict[str, type[Scaler]] = {}
-    """Class variable for all subclasses"""
-
-    def __init__(self, searchspace: pd.DataFrame):
-        self.searchspace = searchspace
-        self.fitted = False
-        self.scale_x: _ScaleFun
-        self.scale_y: _ScaleFun
-        self.unscale_x: _ScaleFun
-        self.unscale_y: _ScaleFun
-        self.unscale_m: _ScaleFun
-        self.unscale_s: _ScaleFun
-
-    @abstractmethod
-    def fit_transform(self, x: Tensor, y: Tensor) -> tuple[Tensor, Tensor]:
-        """Fit the scaler using the given training data and transform the data.
-
-        Args:
-            x: The x-data that should be used.
-            y: The y-data that should be used.
-
-        Returns:
-            The transformed data.
-        """
-
-    def transform(self, x: Tensor) -> Tensor:
-        """Scale a given input.
-
-        Args:
-            x: The given input.
-
-        Returns:
-            The scaled input.
-
-        Raises:
-            RuntimeError: If the scaler is not fitted first.
-        """
-        if not self.fitted:
-            raise RuntimeError("Scaler object must be fitted first.")
-        return self.scale_x(x)
-
-    def untransform(self, mean: Tensor, variance: Tensor) -> tuple[Tensor, Tensor]:
-        """Transform mean values and variances back to the original domain.
-
-        Args:
-            mean: The given mean values.
-            variance: The given variances.
-
-        Returns:
-            The "un-transformed" means and variances.
-
-        Raises:
-            RuntimeError: If the scaler object is not fitted first.
-        """
-        if not self.fitted:
-            raise RuntimeError("Scaler object must be fitted first.")
-        return self.unscale_m(mean), self.unscale_s(variance)
-
-    @classmethod
-    def __init_subclass__(cls, **kwargs):
-        """Register new subclasses dynamically."""
-        super().__init_subclass__(**kwargs)
-        cls.SUBCLASSES[cls.type] = cls
-
-
-class DefaultScaler(Scaler):
-    """A scaler that normalizes inputs to the unit cube and standardizes targets."""
-
-    type = "DEFAULT"
-    # See base class.
-
-    def fit_transform(  # noqa: D102
-        self, x: Tensor, y: Tensor
-    ) -> tuple[Tensor, Tensor]:
-        # See base class.
-
-        import torch
-
-        # Get the searchspace boundaries
-        searchspace = to_tensor(self.searchspace)
-        bounds = torch.vstack(
-            [torch.min(searchspace, dim=0)[0], torch.max(searchspace, dim=0)[0]]
-        )
-
-        # Compute the mean and standard deviation of the training targets
-        mean = torch.mean(y, dim=0)
-        std = torch.std(y, dim=0)
-
-        # Functions for input and target scaling
-        self.scale_x = lambda x: (x - bounds[0]) / (bounds[1] - bounds[0])
-        self.scale_y = lambda x: (x - mean) / std
-
-        # Functions for inverse input and target scaling
-        self.unscale_x = lambda x: x * (bounds[1] - bounds[0]) + bounds[0]
-        self.unscale_y = lambda x: x * std + mean
-
-        # Functions for inverse mean and variance scaling
-        self.unscale_m = lambda x: x * std + mean
-        self.unscale_s = lambda x: x * std**2
-
-        # Flag that the scaler has been fitted
-        self.fitted = True
-
-        return self.scale_x(x), self.scale_y(y)
diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py
index c84f32ab5..d523285b9 100644
--- a/baybe/surrogates/linear.py
+++ b/baybe/surrogates/linear.py
@@ -14,7 +14,7 @@
 from sklearn.linear_model import ARDRegression
 
 from baybe.surrogates.base import GaussianSurrogate
-from baybe.surrogates.utils import autoscale, batchify, catch_constant_targets
+from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
 
 if TYPE_CHECKING:
@@ -22,7 +22,6 @@
 
 
 @catch_constant_targets
-@autoscale
 @define
 class BayesianLinearSurrogate(GaussianSurrogate):
     """A Bayesian linear regression surrogate model."""
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 4a64f82db..303e7e9ec 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -14,7 +14,7 @@
 from ngboost import NGBRegressor
 
 from baybe.surrogates.base import GaussianSurrogate
-from baybe.surrogates.utils import autoscale, batchify, catch_constant_targets
+from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
 
 if TYPE_CHECKING:
@@ -22,7 +22,6 @@
 
 
 @catch_constant_targets
-@autoscale
 @define
 class NGBoostSurrogate(GaussianSurrogate):
     """A natural-gradient-boosting surrogate model."""
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index e4064b985..df19583e3 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -15,7 +15,7 @@
 from sklearn.ensemble import RandomForestRegressor
 
 from baybe.surrogates.base import GaussianSurrogate
-from baybe.surrogates.utils import autoscale, batchify, catch_constant_targets
+from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
 
 if TYPE_CHECKING:
@@ -23,7 +23,6 @@
 
 
 @catch_constant_targets
-@autoscale
 @define
 class RandomForestSurrogate(GaussianSurrogate):
     """A random forest surrogate model."""
diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index b95752ac5..6958c36e4 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -6,9 +6,6 @@
 from functools import wraps
 from typing import TYPE_CHECKING, Any
 
-from baybe.scaler import DefaultScaler
-from baybe.searchspace import SearchSpace
-
 if TYPE_CHECKING:
     from botorch.posteriors import Posterior
     from torch import Tensor
@@ -82,64 +79,6 @@ def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
     return cls
 
 
-def autoscale(cls: type[Surrogate]):
-    """Make a ``Surrogate`` class automatically scale the domain it operates on.
-
-    More specifically, the modified class transforms its inputs before processing them
-    and untransforms the results before returning them. The fitted scaler used for these
-    transformations is stored in the class' objects as an additional temporary
-    attribute.
-
-    Args:
-        cls: The :class:`baybe.surrogates.base.Surrogate` to be augmented.
-
-    Raises:
-        ValueError: If the class already contains an attribute with the same name
-            as the temporary attribute to be added.
-
-    Returns:
-        The modified class.
-    """
-    # Name of the attribute added to store the scaler
-    injected_scaler_attr_name = "_autoscaler"
-
-    if injected_scaler_attr_name in (attr.name for attr in cls.__attrs_attrs__):
-        raise ValueError(
-            f"Cannot apply '{autoscale.__name__}' because "
-            f"'{cls.__name__}' already has an attribute '{injected_scaler_attr_name}' "
-            f"defined."
-        )
-
-    # References to original methods
-    _fit_original = cls._fit
-    _posterior_original = cls._posterior
-
-    def _posterior_new(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
-        scaled = getattr(self, injected_scaler_attr_name).transform(candidates)
-        mean, covar = _posterior_original(self, scaled)
-        return getattr(self, injected_scaler_attr_name).untransform(mean, covar)
-
-    def _fit_new(
-        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> None:
-        scaler = DefaultScaler(searchspace.discrete.comp_rep)
-        scaled_x, scaled_y = scaler.fit_transform(train_x, train_y)
-        try:
-            setattr(self, injected_scaler_attr_name, scaler)
-        except AttributeError as ex:
-            raise TypeError(
-                f"'{autoscale.__name__}' is only applicable to "
-                f"non-slotted classes but '{cls.__name__}' is a slotted class."
-            ) from ex
-        _fit_original(self, searchspace, scaled_x, scaled_y)
-
-    # Replace the methods
-    cls._posterior = _posterior_new
-    cls._fit = _fit_new
-
-    return cls
-
-
 def batchify(
     posterior: Callable[[Surrogate, Tensor], tuple[Tensor, Tensor]],
 ) -> Callable[[Surrogate, Tensor], tuple[Tensor, Tensor]]:
diff --git a/mypy.ini b/mypy.ini
index 67c168b66..10b678f24 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -11,7 +11,6 @@ exclude = (?x)(
           | baybe/deprecation.py
           | baybe/exceptions.py
     	  | baybe/recommenders/naive.py
-          | baybe/scaler.py
           | baybe/simulation.py
           )
 

From 369da451da0e22de948d86f5e89bdf6000902299 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 16 Jul 2024 20:47:34 +0200
Subject: [PATCH 17/97] Make to_tensor also handle numpy arrays

Preparation for use with sklearn's ColumnTransformer, which spits out arrays
---
 baybe/utils/dataframe.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/baybe/utils/dataframe.py b/baybe/utils/dataframe.py
index 5bc09c270..68c408150 100644
--- a/baybe/utils/dataframe.py
+++ b/baybe/utils/dataframe.py
@@ -27,23 +27,23 @@
 
 
 @overload
-def to_tensor(df: pd.DataFrame) -> Tensor:
+def to_tensor(x: np.ndarray | pd.DataFrame, /) -> Tensor:
     ...
 
 
 @overload
-def to_tensor(*dfs: pd.DataFrame) -> Iterator[Tensor]:
+def to_tensor(*x: np.ndarray | pd.DataFrame) -> Iterator[Tensor]:
     ...
 
 
-def to_tensor(*dfs: pd.DataFrame) -> Tensor | Iterator[Tensor]:
-    """Convert a given set of dataframes into tensors (dropping all indices).
+def to_tensor(*x: np.ndarray | pd.DataFrame) -> Tensor | Iterator[Tensor]:
+    """Convert numpy arrays and pandas dataframes into tensors.
 
     Args:
-        *dfs: A set of dataframes
+        *x: The array(s)/dataframe(s) to be converted.
 
     Returns:
-        The provided dataframe(s), transformed into Tensor(s)
+        The provided array(s)/dataframe(s) represented as tensor(s).
     """
     # FIXME This function seems to trigger a problem when some columns in either of
     #  the dfs have a dtype other than int or float (e.g. object, bool). This can
@@ -57,10 +57,12 @@ def to_tensor(*dfs: pd.DataFrame) -> Tensor | Iterator[Tensor]:
     from baybe.utils.torch import DTypeFloatTorch
 
     out = (
-        torch.from_numpy(df.values.astype(DTypeFloatNumpy)).to(DTypeFloatTorch)
-        for df in dfs
+        torch.from_numpy(
+            (xi.values if isinstance(xi, pd.DataFrame) else xi).astype(DTypeFloatNumpy)
+        ).to(DTypeFloatTorch)
+        for xi in x
     )
-    if len(dfs) == 1:
+    if len(x) == 1:
         out = next(out)
     return out
 

From 0ede1ccb74b60028549c6bbd1d450ab530b4729d Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 16 Jul 2024 20:51:51 +0200
Subject: [PATCH 18/97] Replace param_bounds_comp with comp_rep_bounds

---
 CHANGELOG.md                                |  2 ++
 baybe/recommenders/pure/bayesian/botorch.py |  4 ++--
 baybe/searchspace/continuous.py             | 18 +++++++++---------
 baybe/searchspace/core.py                   |  9 ++++-----
 baybe/searchspace/discrete.py               | 20 +++-----------------
 baybe/surrogates/gaussian_process/core.py   |  2 +-
 tests/test_continuous.py                    |  4 ++--
 tests/test_searchspace.py                   |  6 +++---
 8 files changed, 26 insertions(+), 39 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2ea783d7a..9220c2f4b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `Surrogate` models now operate on dataframes in experimental representation instead of
   tensors in computational representation
 - `Surrogate.posterior` models now returns a `Posterior` object
+- `param_bounds_comp` of `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous` has
+  been replaced with `comp_rep_bounds`, which returns a dataframe
 
 ### Added
 - `Surrogate` base class now exposes a `to_botorch` method
diff --git a/baybe/recommenders/pure/bayesian/botorch.py b/baybe/recommenders/pure/bayesian/botorch.py
index 4587789c4..f28b417bb 100644
--- a/baybe/recommenders/pure/bayesian/botorch.py
+++ b/baybe/recommenders/pure/bayesian/botorch.py
@@ -156,7 +156,7 @@ def _recommend_continuous(
 
         points, _ = optimize_acqf(
             acq_function=self._botorch_acqf,
-            bounds=torch.from_numpy(subspace_continuous.param_bounds_comp),
+            bounds=torch.from_numpy(subspace_continuous.comp_rep_bounds.values),
             q=batch_size,
             num_restarts=5,  # TODO make choice for num_restarts
             raw_samples=10,  # TODO make choice for raw_samples
@@ -244,7 +244,7 @@ def _recommend_hybrid(
         # Actual call of the BoTorch optimization routine
         points, _ = optimize_acqf_mixed(
             acq_function=self._botorch_acqf,
-            bounds=torch.from_numpy(searchspace.param_bounds_comp),
+            bounds=torch.from_numpy(searchspace.comp_rep_bounds.values),
             q=batch_size,
             num_restarts=5,  # TODO make choice for num_restarts
             raw_samples=10,  # TODO make choice for raw_samples
diff --git a/baybe/searchspace/continuous.py b/baybe/searchspace/continuous.py
index ab9c2b715..d4fcad744 100644
--- a/baybe/searchspace/continuous.py
+++ b/baybe/searchspace/continuous.py
@@ -30,7 +30,6 @@
 from baybe.serialization import SerialMixin, converter, select_constructor_hook
 from baybe.utils.basic import to_tuple
 from baybe.utils.dataframe import pretty_print_df
-from baybe.utils.numerical import DTypeFloatNumpy
 
 if TYPE_CHECKING:
     from baybe.searchspace.core import SearchSpace
@@ -211,11 +210,12 @@ def param_names(self) -> tuple[str, ...]:
         return tuple(p.name for p in self.parameters)
 
     @property
-    def param_bounds_comp(self) -> np.ndarray:
-        """Return bounds as numpy array."""
-        if not self.parameters:
-            return np.empty((2, 0), dtype=DTypeFloatNumpy)
-        return np.stack([p.bounds.to_ndarray() for p in self.parameters]).T
+    def comp_rep_bounds(self) -> pd.DataFrame:
+        """The minimum and maximum values of the computational representation."""
+        return pd.DataFrame(
+            {p.name: p.bounds.to_tuple() for p in self.parameters},
+            index=["min", "max"],
+        )
 
     def _drop_parameters(self, parameter_names: Collection[str]) -> SubspaceContinuous:
         """Create a copy of the subspace with certain parameters removed.
@@ -324,10 +324,10 @@ def sample_uniform(self, batch_size: int = 1) -> pd.DataFrame:
             and len(self.constraints_lin_ineq) == 0
             and len(self.constraints_cardinality) == 0
         ):
-            return self._sample_from_bounds(batch_size, self.param_bounds_comp)
+            return self._sample_from_bounds(batch_size, self.comp_rep_bounds.values)
 
         if len(self.constraints_cardinality) == 0:
-            return self._sample_from_polytope(batch_size, self.param_bounds_comp)
+            return self._sample_from_polytope(batch_size, self.comp_rep_bounds.values)
 
         return self._sample_from_polytope_with_cardinality_constraints(batch_size)
 
@@ -453,7 +453,7 @@ def sample_from_full_factorial(self, batch_size: int = 1) -> pd.DataFrame:
     def full_factorial(self) -> pd.DataFrame:
         """Get the full factorial of the continuous space."""
         index = pd.MultiIndex.from_product(
-            self.param_bounds_comp.T.tolist(), names=self.param_names
+            self.comp_rep_bounds.values.T.tolist(), names=self.param_names
         )
 
         return pd.DataFrame(index=index).reset_index()
diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py
index 17f1f49f9..d7e6163db 100644
--- a/baybe/searchspace/core.py
+++ b/baybe/searchspace/core.py
@@ -7,7 +7,6 @@
 from enum import Enum
 from typing import cast
 
-import numpy as np
 import pandas as pd
 from attr import define, field
 
@@ -244,10 +243,10 @@ def contains_rdkit(self) -> bool:
         )
 
     @property
-    def param_bounds_comp(self) -> np.ndarray:
-        """Return bounds as tensor."""
-        return np.hstack(
-            [self.discrete.param_bounds_comp, self.continuous.param_bounds_comp]
+    def comp_rep_bounds(self) -> pd.DataFrame:
+        """The minimum and maximum values of the computational representation."""
+        return pd.concat(
+            [self.discrete.comp_rep_bounds, self.continuous.comp_rep_bounds], axis=1
         )
 
     @property
diff --git a/baybe/searchspace/discrete.py b/baybe/searchspace/discrete.py
index 9d41f2f14..c02634432 100644
--- a/baybe/searchspace/discrete.py
+++ b/baybe/searchspace/discrete.py
@@ -537,23 +537,9 @@ def is_empty(self) -> bool:
         return len(self.parameters) == 0
 
     @property
-    def param_bounds_comp(self) -> np.ndarray:
-        """Return bounds as tensor.
-
-        Take bounds from the parameter definitions, but discards bounds belonging to
-        columns that were filtered out during the creation of the space.
-        """
-        if not self.parameters:
-            return np.empty((2, 0))
-        bounds = np.hstack(
-            [
-                np.vstack([p.comp_df[col].min(), p.comp_df[col].max()])
-                for p in self.parameters
-                for col in p.comp_df
-                if col in self.comp_rep.columns
-            ]
-        )
-        return bounds
+    def comp_rep_bounds(self) -> pd.DataFrame:
+        """The minimum and maximum values of the computational representation."""
+        return pd.DataFrame({"min": self.comp_rep.min(), "max": self.comp_rep.max()}).T
 
     @staticmethod
     def estimate_product_space_size(
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index cfd5b99bb..6a2a2afa6 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -62,7 +62,7 @@ def parameter_bounds(self) -> Tensor:
         """Get the search space parameter bounds in BoTorch Format."""
         import torch
 
-        return torch.from_numpy(self.searchspace.param_bounds_comp)
+        return torch.from_numpy(self.searchspace.comp_rep_bounds.values)
 
     def get_numerical_indices(self, n_inputs: int) -> list[int]:
         """Get the indices of the regular numerical model inputs."""
diff --git a/tests/test_continuous.py b/tests/test_continuous.py
index 1f2945715..9ca2816e1 100644
--- a/tests/test_continuous.py
+++ b/tests/test_continuous.py
@@ -13,9 +13,9 @@
 )
 def test_valid_configs(campaign):
     """Test whether the given settings work without error."""
-    print(campaign.searchspace.continuous.param_bounds_comp.flatten())
+    print(campaign.searchspace.continuous.comp_rep_bounds.values.flatten())
 
     assert all(
         np.issubdtype(type(itm), np.floating)
-        for itm in campaign.searchspace.continuous.param_bounds_comp.flatten()
+        for itm in campaign.searchspace.continuous.comp_rep_bounds.values.flatten()
     )
diff --git a/tests/test_searchspace.py b/tests/test_searchspace.py
index 71190f83d..81c083b14 100644
--- a/tests/test_searchspace.py
+++ b/tests/test_searchspace.py
@@ -45,7 +45,7 @@ def test_bounds_order():
     searchspace = SearchSpace.from_product(parameters=parameters)
     expected = np.array([[1.0, 7.0, 4.0, 10.0], [3.0, 9.0, 6.0, 12.0]])
     assert np.array_equal(
-        searchspace.param_bounds_comp,
+        searchspace.comp_rep_bounds.values,
         expected,
     )
 
@@ -59,8 +59,8 @@ def test_empty_parameter_bounds():
     searchspace_discrete = SubspaceDiscrete.from_product(parameters=parameters)
     searchspace_continuous = SubspaceContinuous(parameters=parameters)
     expected = np.empty((2, 0))
-    assert np.array_equal(searchspace_discrete.param_bounds_comp, expected)
-    assert np.array_equal(searchspace_continuous.param_bounds_comp, expected)
+    assert np.array_equal(searchspace_discrete.comp_rep_bounds.values, expected)
+    assert np.array_equal(searchspace_continuous.comp_rep_bounds.values, expected)
 
 
 def test_discrete_searchspace_creation_from_dataframe():

From 00c40ae50f5df5ab8fede77be2e1ec3b826fe075 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 17 Jul 2024 09:53:14 +0200
Subject: [PATCH 19/97] Draft input scaling mechanism

---
 baybe/surrogates/base.py                  | 36 ++++++++++++++++++++--
 baybe/surrogates/gaussian_process/core.py |  9 ++++++
 baybe/surrogates/ngboost.py               | 10 ++++++
 baybe/surrogates/random_forest.py         | 10 ++++++
 baybe/utils/scaling.py                    | 37 +++++++++++++++++++++++
 5 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 baybe/utils/scaling.py

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index aa7ffc582..2b7029a21 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -19,6 +19,7 @@
 
 from baybe.exceptions import ModelNotTrainedError
 from baybe.objectives.base import Objective
+from baybe.parameters.base import Parameter
 from baybe.searchspace import SearchSpace
 from baybe.serialization.core import (
     converter,
@@ -27,13 +28,14 @@
 )
 from baybe.serialization.mixin import SerialMixin
 from baybe.utils.dataframe import to_tensor
+from baybe.utils.scaling import ScalingMethod, make_scaler
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
     from botorch.posteriors import GPyTorchPosterior, Posterior
+    from sklearn.compose import ColumnTransformer
     from torch import Tensor
 
-
 _ONNX_ENCODING = "latin-1"
 """Constant signifying the encoding for onnx byte strings in pretrained models.
 
@@ -83,6 +85,32 @@ def to_botorch(self) -> Model:
 
         return AdapterModel(self)
 
+    @staticmethod
+    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
+        """Return the scaling method to be used for the given parameter."""
+        return ScalingMethod.MINMAX
+
+    def _make_input_scaler(
+        self, searchspace: SearchSpace, measurements: pd.DataFrame
+    ) -> ColumnTransformer:
+        """Make a scaler to be used for transforming computational dataframes."""
+        from sklearn.compose import make_column_transformer
+
+        # Create the composite scaler from the parameter-wise scaler objects
+        # TODO: Filter down to columns that actually remain in the comp rep of the
+        #   searchspace, since the transformer can break down otherwise.
+        transformers = [
+            (make_scaler(self._get_parameter_scaling(p)), p.comp_df.columns)
+            for p in searchspace.parameters
+        ]
+        scaler = make_column_transformer(*transformers)
+
+        # TODO: Decide whether scaler is to be fit to parameter bounds and/or
+        #   extreme points in the given measurement data
+        scaler.fit(searchspace.comp_rep_bounds)
+
+        return scaler
+
     def transform_inputs(self, data: pd.DataFrame) -> pd.DataFrame:
         """Transform an experimental parameter dataframe."""
         if self._input_transform is None:
@@ -148,8 +176,12 @@ def fit(
                 "Continuous search spaces are currently only supported by GPs."
             )
 
+        input_scaler = self._make_input_scaler(searchspace, measurements)
+
         # Store context-specific transformations
-        self._input_transform = lambda x: searchspace.transform(x, allow_missing=True)
+        self._input_transform = lambda x: input_scaler.transform(
+            searchspace.transform(x, allow_missing=True)
+        )
         self._target_transform = lambda x: objective.transform(x)
 
         # Transform and fit
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 6a2a2afa6..e612ce506 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -8,6 +8,7 @@
 from attrs.validators import instance_of
 
 from baybe.objective import Objective
+from baybe.parameters.base import Parameter
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.base import Surrogate
 from baybe.surrogates.gaussian_process.kernel_factory import (
@@ -22,6 +23,7 @@
     DefaultKernelFactory,
     _default_noise_factory,
 )
+from baybe.utils.scaling import ScalingMethod
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
@@ -108,6 +110,13 @@ def to_botorch(self) -> Model:  # noqa: D102
 
         return self._model
 
+    @staticmethod
+    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
+        # See base class.
+
+        # For GPs, we use botorch's built-in machinery for scaling
+        return ScalingMethod.IDENTITY
+
     @staticmethod
     def _get_model_context(
         searchspace: SearchSpace, objective: Objective
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 303e7e9ec..943d5bbf5 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -6,6 +6,7 @@
 Since we plan to refactor the surrogates, this part of the documentation will be
 available in the future. Thus, please have a look in the source code directly.
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, ClassVar
@@ -13,9 +14,11 @@
 from attr import define, field
 from ngboost import NGBRegressor
 
+from baybe.parameters.base import Parameter
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
+from baybe.utils.scaling import ScalingMethod
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -50,6 +53,13 @@ class NGBoostSurrogate(GaussianSurrogate):
     def __attrs_post_init__(self):
         self.model_params = {**self._default_model_params, **self.model_params}
 
+    @staticmethod
+    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
+        # See base class.
+
+        # Tree-like models do not require any input scaling
+        return ScalingMethod.IDENTITY
+
     @batchify
     def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         # See base class.
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index df19583e3..eb06f65d9 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -6,6 +6,7 @@
 Since we plan to refactor the surrogates, this part of the documentation will be
 available in the future. Thus, please have a look in the source code directly.
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, ClassVar
@@ -14,9 +15,11 @@
 from attr import define, field
 from sklearn.ensemble import RandomForestRegressor
 
+from baybe.parameters.base import Parameter
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
+from baybe.utils.scaling import ScalingMethod
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -45,6 +48,13 @@ class RandomForestSurrogate(GaussianSurrogate):
     _model: RandomForestRegressor | None = field(init=False, default=None, eq=False)
     """The actual model."""
 
+    @staticmethod
+    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
+        # See base class.
+
+        # Tree-like models do not require any input scaling
+        return ScalingMethod.IDENTITY
+
     @batchify
     def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
         # See base class.
diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
new file mode 100644
index 000000000..bc4455d48
--- /dev/null
+++ b/baybe/utils/scaling.py
@@ -0,0 +1,37 @@
+"""Scaling utilities."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import TYPE_CHECKING, Literal, TypeAlias
+
+if TYPE_CHECKING:
+    from sklearn.base import BaseEstimator, TransformerMixin
+
+    Scaler: TypeAlias = BaseEstimator | TransformerMixin
+
+
+class ScalingMethod(Enum):
+    """Available scaling methods."""
+
+    IDENTITY = "IDENTITY"
+    """Identity transformation (no scaling applied)."""
+
+    MINMAX = "MINMAX"
+    """Min-max scaling, mapping the observed value range to [0, 1]."""
+
+    MAXABS = "MAXABS"
+    """Max-abs scaling, normalizing by the largest observed absolute."""
+
+
+def make_scaler(method: ScalingMethod, /) -> Scaler | Literal["passthrough"]:
+    """Create a scaler object based on the specified method."""
+    from sklearn.preprocessing import MaxAbsScaler, MinMaxScaler
+
+    match method:
+        case ScalingMethod.IDENTITY:
+            return "passthrough"
+        case ScalingMethod.MINMAX:
+            return MinMaxScaler()
+        case ScalingMethod.MAXABS:
+            return MaxAbsScaler()

From 79f8f4461ab2431c7c10e55176d8a27eba18873f Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 19 Jul 2024 09:24:41 +0200
Subject: [PATCH 20/97] Introduce ScalerProtocol class

---
 baybe/surrogates/base.py                  | 19 ++++++-----
 baybe/surrogates/gaussian_process/core.py | 10 +++---
 baybe/surrogates/ngboost.py               | 10 +++---
 baybe/surrogates/random_forest.py         | 10 +++---
 baybe/utils/scaling.py                    | 39 +++++++----------------
 5 files changed, 41 insertions(+), 47 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 2b7029a21..743862cfb 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -4,7 +4,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 import pandas as pd
 from attrs import define, field
@@ -16,6 +16,7 @@
     UnstructuredValue,
     UnstructureHook,
 )
+from sklearn.preprocessing import MinMaxScaler
 
 from baybe.exceptions import ModelNotTrainedError
 from baybe.objectives.base import Objective
@@ -28,7 +29,7 @@
 )
 from baybe.serialization.mixin import SerialMixin
 from baybe.utils.dataframe import to_tensor
-from baybe.utils.scaling import ScalingMethod, make_scaler
+from baybe.utils.scaling import ScalerProtocol
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
@@ -86,9 +87,11 @@ def to_botorch(self) -> Model:
         return AdapterModel(self)
 
     @staticmethod
-    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
-        """Return the scaling method to be used for the given parameter."""
-        return ScalingMethod.MINMAX
+    def _get_parameter_scaler(
+        parameter: Parameter,
+    ) -> ScalerProtocol | Literal["passthrough"]:
+        """Return the scaler to be used for the given parameter."""
+        return MinMaxScaler()
 
     def _make_input_scaler(
         self, searchspace: SearchSpace, measurements: pd.DataFrame
@@ -100,7 +103,7 @@ def _make_input_scaler(
         # TODO: Filter down to columns that actually remain in the comp rep of the
         #   searchspace, since the transformer can break down otherwise.
         transformers = [
-            (make_scaler(self._get_parameter_scaling(p)), p.comp_df.columns)
+            (self._get_parameter_scaler(p), p.comp_df.columns)
             for p in searchspace.parameters
         ]
         scaler = make_column_transformer(*transformers)
@@ -225,7 +228,7 @@ def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
 
 
 def _make_hook_decode_onnx_str(
-    raw_unstructure_hook: UnstructureHook
+    raw_unstructure_hook: UnstructureHook,
 ) -> UnstructureHook:
     """Wrap an unstructuring hook to let it also decode the contained ONNX string."""
 
@@ -253,7 +256,7 @@ def wrapper(dct: UnstructuredValue, _: TargetType) -> StructuredValue:
 
 
 def _block_serialize_custom_architecture(
-    raw_unstructure_hook: UnstructureHook
+    raw_unstructure_hook: UnstructureHook,
 ) -> UnstructureHook:
     """Raise error if attempt to serialize a custom architecture surrogate."""
     # TODO: Ideally, this hook should be removed and unstructuring the Surrogate
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index e612ce506..e83c984c9 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, ClassVar, Literal
 
 from attrs import define, field
 from attrs.validators import instance_of
@@ -23,7 +23,7 @@
     DefaultKernelFactory,
     _default_noise_factory,
 )
-from baybe.utils.scaling import ScalingMethod
+from baybe.utils.scaling import ScalerProtocol
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
@@ -111,11 +111,13 @@ def to_botorch(self) -> Model:  # noqa: D102
         return self._model
 
     @staticmethod
-    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
+    def _get_parameter_scaler(
+        parameter: Parameter,
+    ) -> ScalerProtocol | Literal["passthrough"]:
         # See base class.
 
         # For GPs, we use botorch's built-in machinery for scaling
-        return ScalingMethod.IDENTITY
+        return "passthrough"
 
     @staticmethod
     def _get_model_context(
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 943d5bbf5..937caa27a 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 from attr import define, field
 from ngboost import NGBRegressor
@@ -18,7 +18,7 @@
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
-from baybe.utils.scaling import ScalingMethod
+from baybe.utils.scaling import ScalerProtocol
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -54,11 +54,13 @@ def __attrs_post_init__(self):
         self.model_params = {**self._default_model_params, **self.model_params}
 
     @staticmethod
-    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
+    def _get_parameter_scaler(
+        parameter: Parameter,
+    ) -> ScalerProtocol | Literal["passthrough"]:
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return ScalingMethod.IDENTITY
+        return "passthrough"
 
     @batchify
     def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index eb06f65d9..6ef9e0196 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 import numpy as np
 from attr import define, field
@@ -19,7 +19,7 @@
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
-from baybe.utils.scaling import ScalingMethod
+from baybe.utils.scaling import ScalerProtocol
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -49,11 +49,13 @@ class RandomForestSurrogate(GaussianSurrogate):
     """The actual model."""
 
     @staticmethod
-    def _get_parameter_scaling(parameter: Parameter) -> ScalingMethod:
+    def _get_parameter_scaler(
+        parameter: Parameter,
+    ) -> ScalerProtocol | Literal["passthrough"]:
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return ScalingMethod.IDENTITY
+        return "passthrough"
 
     @batchify
     def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index bc4455d48..0bff2f941 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -2,36 +2,21 @@
 
 from __future__ import annotations
 
-from enum import Enum
-from typing import TYPE_CHECKING, Literal, TypeAlias
+from typing import Protocol
 
-if TYPE_CHECKING:
-    from sklearn.base import BaseEstimator, TransformerMixin
+import pandas as pd
 
-    Scaler: TypeAlias = BaseEstimator | TransformerMixin
 
+class ScalerProtocol(Protocol):
+    """Type protocol specifying the interface scalers need to implement.
 
-class ScalingMethod(Enum):
-    """Available scaling methods."""
+    The protocol is compatible with sklearn scalers such as
+    :class:`sklearn.preprocessing.MinMaxScaler` or
+    :class:`sklearn.preprocessing.MaxAbsScaler`.
+    """
 
-    IDENTITY = "IDENTITY"
-    """Identity transformation (no scaling applied)."""
+    def fit(self, df: pd.DataFrame, /) -> None:
+        """Fit the scaler to a given data set."""
 
-    MINMAX = "MINMAX"
-    """Min-max scaling, mapping the observed value range to [0, 1]."""
-
-    MAXABS = "MAXABS"
-    """Max-abs scaling, normalizing by the largest observed absolute."""
-
-
-def make_scaler(method: ScalingMethod, /) -> Scaler | Literal["passthrough"]:
-    """Create a scaler object based on the specified method."""
-    from sklearn.preprocessing import MaxAbsScaler, MinMaxScaler
-
-    match method:
-        case ScalingMethod.IDENTITY:
-            return "passthrough"
-        case ScalingMethod.MINMAX:
-            return MinMaxScaler()
-        case ScalingMethod.MAXABS:
-            return MaxAbsScaler()
+    def transform(self, df: pd.DataFrame, /) -> pd.DataFrame:
+        """Transform a data using the fitted scaling logic."""

From 24f2c490beca9ec8501c7a73d63daccf4b7c4a09 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 19 Jul 2024 09:35:57 +0200
Subject: [PATCH 21/97] Make transformation return a dataframe

---
 baybe/surrogates/base.py | 13 +++++++++----
 ruff.toml                |  3 +++
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 743862cfb..8e8e90329 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -106,7 +106,7 @@ def _make_input_scaler(
             (self._get_parameter_scaler(p), p.comp_df.columns)
             for p in searchspace.parameters
         ]
-        scaler = make_column_transformer(*transformers)
+        scaler = make_column_transformer(*transformers, verbose_feature_names_out=False)
 
         # TODO: Decide whether scaler is to be fit to parameter bounds and/or
         #   extreme points in the given measurement data
@@ -181,10 +181,15 @@ def fit(
 
         input_scaler = self._make_input_scaler(searchspace, measurements)
 
+        def transform_inputs(df: pd.DataFrame, /) -> pd.DataFrame:
+            """Fitted input transformation pipeline."""
+            out = input_scaler.transform(searchspace.transform(df, allow_missing=True))
+            return pd.DataFrame(
+                out, index=df.index, columns=input_scaler.get_feature_names_out()
+            )
+
         # Store context-specific transformations
-        self._input_transform = lambda x: input_scaler.transform(
-            searchspace.transform(x, allow_missing=True)
-        )
+        self._input_transform = transform_inputs
         self._target_transform = lambda x: objective.transform(x)
 
         # Transform and fit
diff --git a/ruff.toml b/ruff.toml
index 6add2d104..e7278f12b 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -53,3 +53,6 @@ select=[
 
 [lint.pydocstyle]
 convention = "google"
+
+[lint.isort]
+known-third-party = ["streamlit"]  # needed because of our "streamlit" folder

From 2938c488a2e84cf54a1eb1ba3feeb18e18beef8a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 19 Jul 2024 09:39:04 +0200
Subject: [PATCH 22/97] Update streamlit dev script

---
 streamlit/surrogate_models.py | 146 +++++++++++++++++-----------------
 1 file changed, 75 insertions(+), 71 deletions(-)

diff --git a/streamlit/surrogate_models.py b/streamlit/surrogate_models.py
index 00027580b..a699b1f59 100644
--- a/streamlit/surrogate_models.py
+++ b/streamlit/surrogate_models.py
@@ -1,28 +1,29 @@
-"""Compare different surrogate model and verify their predictions.
+"""# Surrogate Models
 
 The purpose of this script is to compare different surrogate models and verify that
-their predictions are invariant to changes in scale of the underlying target function.
+their predictions are invariant to changes in location/scale of the underlying target
+function.
 
-This means that the shown function approximation should always appear visually the same
-when the input and output scales are changed.
-"""
+This means that the displayed function approximation should always look the same when
+the input and output locations/scales are changed.
+"""  # noqa: D415
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-import torch
-from botorch.optim import optimize_acqf_discrete
+import streamlit as st
 from funcy import rpartial
 
-import streamlit as st
-from baybe.acquisition import qExpectedImprovement
-from baybe.parameters import NumericalDiscreteParameter
+from baybe.parameters.numerical import NumericalDiscreteParameter
+from baybe.recommenders.pure.bayesian.botorch import BotorchRecommender
 from baybe.searchspace import SearchSpace
 from baybe.surrogates import CustomONNXSurrogate
 from baybe.surrogates.base import Surrogate
+from baybe.targets.numerical import NumericalTarget
 from baybe.utils.basic import get_subclasses
+from baybe.utils.random import set_random_seed
 
-# define constants
+# Number of values used for the input parameter
 N_PARAMETER_VALUES = 1000
 
 
@@ -60,13 +61,10 @@ def linear(
 
 def main():
     """Create the streamlit dashboard."""
-    # basic settings
-    plt.style.use("seaborn-v0_8-paper")  # 'seaborn' is depreciated since matplotlib 3.6
-
-    # show docstring in dashboard
+    # Show module docstring in dashboard
     st.info(__doc__)
 
-    # define all available test functions
+    # Collect all available test functions
     test_functions = {
         "Sine": sin,
         "Constant": constant,
@@ -74,89 +72,95 @@ def main():
         "Cubic": cubic,
     }
 
-    # collect all available surrogate models
+    # Collect all available surrogate models
     surrogate_model_classes = {
-        surr.__name__: surr
-        for surr in get_subclasses(Surrogate)
-        if not issubclass(surr, CustomONNXSurrogate)
+        cls.__name__: cls
+        for cls in get_subclasses(Surrogate)
+        if not issubclass(cls, CustomONNXSurrogate)
     }
 
-    # simulation parameters
-    random_seed = int(st.sidebar.number_input("Random seed", value=1337))
-    function_name = st.sidebar.selectbox("Test function", list(test_functions.keys()))
-    surrogate_name = st.sidebar.selectbox(
+    # Streamlit simulation parameters
+    st_random_seed = int(st.sidebar.number_input("Random seed", value=1337))
+    st_function_name = st.sidebar.selectbox(
+        "Test function", list(test_functions.keys())
+    )
+    st_surrogate_name = st.sidebar.selectbox(
         "Surrogate model", list(surrogate_model_classes.keys())
     )
-    n_training_points = st.sidebar.slider("Number of training points", 1, 20, 5)
-    n_recommendations = st.sidebar.slider("Number of recommendations", 1, 20, 5)
+    st_n_training_points = st.sidebar.slider("Number of training points", 1, 20, 5)
+    st_n_recommendations = st.sidebar.slider("Number of recommendations", 1, 20, 5)
     st.sidebar.markdown("---")
     st.sidebar.markdown(
         """
-        The plot should remain static (except for the axis labels) when changing the
-        following parameters.
+        When scaling is implemented correctly, the plot should remain static (except for
+        the axis labels) when changing the following parameters:
         """
     )
-    upper_parameter_limit = st.sidebar.slider("Upper parameter limit", 0.0, 100.0, 1.0)
-    lower_parameter_limit = st.sidebar.slider("Lower parameter limit", -100.0, 0.0, 0.0)
-    function_amplitude = st.sidebar.slider("Function amplitude", 1.0, 100.0, 1.0)
-    function_bias = st.sidebar.slider("Function bias", -100.0, 100.0, 0.0)
+    st_upper_parameter_limit = st.sidebar.slider(
+        "Upper parameter limit", 0.0, 100.0, 1.0
+    )
+    st_lower_parameter_limit = st.sidebar.slider(
+        "Lower parameter limit", -100.0, 0.0, 0.0
+    )
+    st_function_amplitude = st.sidebar.slider("Function amplitude", 1.0, 100.0, 1.0)
+    st_function_bias = st.sidebar.slider("Function bias", -100.0, 100.0, 0.0)
 
-    # fix the chosen random seed
-    np.random.seed(random_seed)
-    torch.manual_seed(random_seed)
+    # Set the chosen random seed
+    set_random_seed(st_random_seed)
 
-    # select the test function and the surrogate model class
+    # Construct the specific test function
     fun = rpartial(
-        test_functions[function_name],
-        lower_parameter_limit,
-        upper_parameter_limit,
-        function_amplitude,
-        function_bias,
+        test_functions[st_function_name],
+        st_lower_parameter_limit,
+        st_upper_parameter_limit,
+        st_function_amplitude,
+        st_function_bias,
     )
 
-    # create the input grid and corresponding target values
-    test_x = torch.linspace(
-        lower_parameter_limit, upper_parameter_limit, N_PARAMETER_VALUES
+    # Create the training data
+    train_x = np.random.uniform(
+        st_lower_parameter_limit, st_upper_parameter_limit, st_n_training_points
     )
-    test_y = torch.from_numpy(fun(test_x.numpy()))
+    train_y = fun(train_x)
+    measurements = pd.DataFrame({"x": train_x, "y": train_y})
 
-    # randomly select the specified number of training data points
-    train_idx = np.random.choice(
-        range(N_PARAMETER_VALUES), n_training_points, replace=False
+    # Create the plotting grid and corresponding target values
+    test_x = np.linspace(
+        st_lower_parameter_limit, st_upper_parameter_limit, N_PARAMETER_VALUES
     )
-    train_x = test_x[train_idx]
-    train_y = test_y[train_idx]
-
-    # create the searchspace object
-    param = NumericalDiscreteParameter(name="param", values=test_x.numpy().tolist())
-    searchspace = SearchSpace.from_product(parameters=[param])
+    test_y = fun(test_x)
+    candidates = pd.DataFrame({"x": test_x, "y": test_y})
+
+    # Create the searchspace and objective
+    parameter = NumericalDiscreteParameter(
+        name="x",
+        values=np.linspace(
+            st_lower_parameter_limit, st_upper_parameter_limit, N_PARAMETER_VALUES
+        ),
+    )
+    searchspace = SearchSpace.from_product(parameters=[parameter])
+    objective = NumericalTarget(name="y", mode="MAX").to_objective()
 
-    # create the surrogate model, train it, and get its predictions
-    surrogate_model = surrogate_model_classes[surrogate_name]()
-    surrogate_model.fit(searchspace, train_x.unsqueeze(-1), train_y.unsqueeze(-1))
+    # Create the surrogate model and the recommender
+    surrogate_model = surrogate_model_classes[st_surrogate_name]()
+    recommender = BotorchRecommender(surrogate_model=surrogate_model)
 
-    # recommend next experiments
-    # TODO: use BayBE recommender and add widgets for recommender selection
-    acqf = qExpectedImprovement().to_botorch(
-        surrogate_model, searchspace, pd.DataFrame(train_x), pd.DataFrame(train_y)
+    # Get the recommendations and extract the posterior mean / standard deviation
+    recommendations = recommender.recommend(
+        st_n_recommendations, searchspace, objective, measurements
     )
-    recommendatations = optimize_acqf_discrete(
-        acqf, q=n_recommendations, choices=test_x.unsqueeze(-1)
-    )[0]
-
-    # create the mean and standard deviation predictions for the entire search space
-    mean, covar = surrogate_model.posterior(test_x.unsqueeze(-1))
-    mean = mean.detach().numpy()
-    std = covar.diag().sqrt().detach().numpy()
+    posterior = surrogate_model.posterior(candidates)
+    mean = posterior.mean.squeeze().detach().numpy()
+    std = posterior.stddev.detach().numpy()
 
-    # visualize the test function, training points, model predictions, recommendations
+    # Visualize the test function, training points, model predictions, recommendations
     fig = plt.figure()
     plt.plot(test_x, test_y, color="tab:blue", label="Test function")
     plt.plot(train_x, train_y, "o", color="tab:blue")
     plt.plot(test_x, mean, color="tab:red", label="Surrogate model")
     plt.fill_between(test_x, mean - std, mean + std, alpha=0.2, color="tab:red")
     plt.vlines(
-        recommendatations, *plt.gca().get_ylim(), color="k", label="Recommendations"
+        recommendations, *plt.gca().get_ylim(), color="k", label="Recommendations"
     )
     plt.legend()
     st.pyplot(fig)

From ae1a36676996f731ad0160bab88c6488d0e5abef Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 19 Jul 2024 12:14:05 +0200
Subject: [PATCH 23/97] Fix handling of dropped columns in ColumnTransformer

---
 CHANGELOG.md                    | 2 ++
 baybe/searchspace/continuous.py | 5 +++++
 baybe/searchspace/core.py       | 5 +++++
 baybe/searchspace/discrete.py   | 5 +++++
 baybe/surrogates/base.py        | 8 +++++---
 5 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9220c2f4b..cdd4cf0fa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,8 @@ _ `_optional` subpackage for managing optional dependencies
 - `transform` methods of `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous`
   now take additional `allow_missing` and `allow_extra` keyword arguments
 - `GaussianSurrogate` base class for surrogate models with Gaussian posteriors
+- `comp_rep_columns` property for `SearchSpace`, `SubspaceDiscrete`
+  and `SubspaceContinuous`
 
 ### Changed
 - Passing an `Objective` to `Campaign` is now optional
diff --git a/baybe/searchspace/continuous.py b/baybe/searchspace/continuous.py
index d4fcad744..76f7212c1 100644
--- a/baybe/searchspace/continuous.py
+++ b/baybe/searchspace/continuous.py
@@ -209,6 +209,11 @@ def param_names(self) -> tuple[str, ...]:
         """Return list of parameter names."""
         return tuple(p.name for p in self.parameters)
 
+    @property
+    def comp_rep_columns(self) -> tuple[str, ...]:
+        """The columns spanning the computational representation."""
+        return self.param_names
+
     @property
     def comp_rep_bounds(self) -> pd.DataFrame:
         """The minimum and maximum values of the computational representation."""
diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py
index d7e6163db..262f35f88 100644
--- a/baybe/searchspace/core.py
+++ b/baybe/searchspace/core.py
@@ -242,6 +242,11 @@ def contains_rdkit(self) -> bool:
             p.encoding is SubstanceEncoding.RDKIT for p in self.discrete.parameters
         )
 
+    @property
+    def comp_rep_columns(self) -> tuple[str, ...]:
+        """The columns spanning the computational representation."""
+        return self.discrete.comp_rep_columns + self.continuous.comp_rep_columns
+
     @property
     def comp_rep_bounds(self) -> pd.DataFrame:
         """The minimum and maximum values of the computational representation."""
diff --git a/baybe/searchspace/discrete.py b/baybe/searchspace/discrete.py
index c02634432..de9553229 100644
--- a/baybe/searchspace/discrete.py
+++ b/baybe/searchspace/discrete.py
@@ -536,6 +536,11 @@ def is_empty(self) -> bool:
         """Return whether this subspace is empty."""
         return len(self.parameters) == 0
 
+    @property
+    def comp_rep_columns(self) -> tuple[str, ...]:
+        """The columns spanning the computational representation."""
+        return tuple(self.comp_rep.columns)
+
     @property
     def comp_rep_bounds(self) -> pd.DataFrame:
         """The minimum and maximum values of the computational representation."""
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 8e8e90329..36b237833 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -100,10 +100,12 @@ def _make_input_scaler(
         from sklearn.compose import make_column_transformer
 
         # Create the composite scaler from the parameter-wise scaler objects
-        # TODO: Filter down to columns that actually remain in the comp rep of the
-        #   searchspace, since the transformer can break down otherwise.
+        # TODO: Fix the parameter comp rep column access for continuous parameters
         transformers = [
-            (self._get_parameter_scaler(p), p.comp_df.columns)
+            (
+                self._get_parameter_scaler(p),
+                [c for c in p.comp_df.columns if c in searchspace.comp_rep_columns],
+            )
             for p in searchspace.parameters
         ]
         scaler = make_column_transformer(*transformers, verbose_feature_names_out=False)

From 50681481d3bd209d93f104bb1be30d3d72a32592 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 19 Jul 2024 12:18:35 +0200
Subject: [PATCH 24/97] Remove obsolete TODO note

---
 baybe/surrogates/base.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 36b237833..87b224b28 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -173,7 +173,6 @@ def fit(
                 f"surrogate model type ({self.__class__.__name__}) does not "
                 f"support transfer learning."
             )
-        # TODO: Adjust scale_model decorator to support other model types as well.
         if (not searchspace.continuous.is_empty) and (
             "GaussianProcess" not in self.__class__.__name__
         ):

From fb149277002b8035c190d42feea076d5b8a1650c Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 19 Jul 2024 14:51:53 +0200
Subject: [PATCH 25/97] Make surrogate scaling work with continuous parameters

---
 CHANGELOG.md                    |  4 ++--
 baybe/parameters/base.py        | 19 +++++++++++++++----
 baybe/parameters/numerical.py   |  5 +++++
 baybe/searchspace/continuous.py |  2 +-
 baybe/searchspace/discrete.py   |  5 ++++-
 baybe/surrogates/base.py        |  3 +--
 6 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cdd4cf0fa..dc4bf2504 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,8 +35,8 @@ _ `_optional` subpackage for managing optional dependencies
 - `transform` methods of `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous`
   now take additional `allow_missing` and `allow_extra` keyword arguments
 - `GaussianSurrogate` base class for surrogate models with Gaussian posteriors
-- `comp_rep_columns` property for `SearchSpace`, `SubspaceDiscrete`
-  and `SubspaceContinuous`
+- `comp_rep_columns` property for `Parameter`, `SearchSpace`, `SubspaceDiscrete`
+  and `SubspaceContinuous` classes
 
 ### Changed
 - Passing an `Objective` to `Campaign` is now optional
diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py
index 809d7eab0..d7e02619f 100644
--- a/baybe/parameters/base.py
+++ b/baybe/parameters/base.py
@@ -48,10 +48,6 @@ def is_in_range(self, item: Any) -> bool:
             ``True`` if the item is within the parameter range, ``False`` otherwise.
         """
 
-    @abstractmethod
-    def summary(self) -> dict:
-        """Return a custom summarization of the parameter."""
-
     def __str__(self) -> str:
         return str(self.summary())
 
@@ -65,6 +61,15 @@ def is_discrete(self) -> bool:
         """Boolean indicating if this is a discrete parameter."""
         return isinstance(self, DiscreteParameter)
 
+    @property
+    @abstractmethod
+    def comp_rep_columns(self) -> tuple[str, ...]:
+        """The columns spanning the computational representation."""
+
+    @abstractmethod
+    def summary(self) -> dict:
+        """Return a custom summarization of the parameter."""
+
 
 @define(frozen=True, slots=False)
 class DiscreteParameter(Parameter, ABC):
@@ -84,8 +89,14 @@ def values(self) -> tuple:
     @cached_property
     @abstractmethod
     def comp_df(self) -> pd.DataFrame:
+        # TODO: Should be renamed to `comp_rep`
         """Return the computational representation of the parameter."""
 
+    @property
+    def comp_rep_columns(self) -> tuple[str, ...]:  # noqa: D102
+        # See base class.
+        return tuple(self.comp_df.columns)
+
     def is_in_range(self, item: Any) -> bool:  # noqa: D102
         # See base class.
         return item in self.values
diff --git a/baybe/parameters/numerical.py b/baybe/parameters/numerical.py
index 93dbb8e7f..4f92043c7 100644
--- a/baybe/parameters/numerical.py
+++ b/baybe/parameters/numerical.py
@@ -132,6 +132,11 @@ def is_in_range(self, item: float) -> bool:  # noqa: D102
 
         return self.bounds.contains(item)
 
+    @property
+    def comp_rep_columns(self) -> tuple[str, ...]:  # noqa: D102
+        # See base class.
+        return (self.name,)
+
     def summary(self) -> dict:  # noqa: D102
         # See base class.
         param_dict = dict(
diff --git a/baybe/searchspace/continuous.py b/baybe/searchspace/continuous.py
index 76f7212c1..f612b24a9 100644
--- a/baybe/searchspace/continuous.py
+++ b/baybe/searchspace/continuous.py
@@ -212,7 +212,7 @@ def param_names(self) -> tuple[str, ...]:
     @property
     def comp_rep_columns(self) -> tuple[str, ...]:
         """The columns spanning the computational representation."""
-        return self.param_names
+        return tuple(chain.from_iterable(p.comp_rep_columns for p in self.parameters))
 
     @property
     def comp_rep_bounds(self) -> pd.DataFrame:
diff --git a/baybe/searchspace/discrete.py b/baybe/searchspace/discrete.py
index de9553229..8e2572c23 100644
--- a/baybe/searchspace/discrete.py
+++ b/baybe/searchspace/discrete.py
@@ -539,6 +539,9 @@ def is_empty(self) -> bool:
     @property
     def comp_rep_columns(self) -> tuple[str, ...]:
         """The columns spanning the computational representation."""
+        # We go via `comp_rep` here instead of using the columns of the individual
+        # parameters because the search space potentially uses only a subset of the
+        # columns due to decorrelation
         return tuple(self.comp_rep.columns)
 
     @property
@@ -548,7 +551,7 @@ def comp_rep_bounds(self) -> pd.DataFrame:
 
     @staticmethod
     def estimate_product_space_size(
-        parameters: Sequence[DiscreteParameter]
+        parameters: Sequence[DiscreteParameter],
     ) -> MemorySize:
         """Estimate an upper bound for the memory size of a product space.
 
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 87b224b28..938a4048e 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -100,11 +100,10 @@ def _make_input_scaler(
         from sklearn.compose import make_column_transformer
 
         # Create the composite scaler from the parameter-wise scaler objects
-        # TODO: Fix the parameter comp rep column access for continuous parameters
         transformers = [
             (
                 self._get_parameter_scaler(p),
-                [c for c in p.comp_df.columns if c in searchspace.comp_rep_columns],
+                [c for c in p.comp_rep_columns if c in searchspace.comp_rep_columns],
             )
             for p in searchspace.parameters
         ]

From c3a4cc631be9d50ba5c2d7175e1b45180bc2a00f Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 19 Jul 2024 15:07:00 +0200
Subject: [PATCH 26/97] Rename _get_parameter_scaler to _make_parameter_scaler

---
 baybe/surrogates/base.py                  | 4 ++--
 baybe/surrogates/gaussian_process/core.py | 2 +-
 baybe/surrogates/ngboost.py               | 2 +-
 baybe/surrogates/random_forest.py         | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 938a4048e..21e268763 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -87,7 +87,7 @@ def to_botorch(self) -> Model:
         return AdapterModel(self)
 
     @staticmethod
-    def _get_parameter_scaler(
+    def _make_parameter_scaler(
         parameter: Parameter,
     ) -> ScalerProtocol | Literal["passthrough"]:
         """Return the scaler to be used for the given parameter."""
@@ -102,7 +102,7 @@ def _make_input_scaler(
         # Create the composite scaler from the parameter-wise scaler objects
         transformers = [
             (
-                self._get_parameter_scaler(p),
+                self._make_parameter_scaler(p),
                 [c for c in p.comp_rep_columns if c in searchspace.comp_rep_columns],
             )
             for p in searchspace.parameters
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index e83c984c9..1ff56c1d1 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -111,7 +111,7 @@ def to_botorch(self) -> Model:  # noqa: D102
         return self._model
 
     @staticmethod
-    def _get_parameter_scaler(
+    def _make_parameter_scaler(
         parameter: Parameter,
     ) -> ScalerProtocol | Literal["passthrough"]:
         # See base class.
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 937caa27a..6a4f98d2f 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -54,7 +54,7 @@ def __attrs_post_init__(self):
         self.model_params = {**self._default_model_params, **self.model_params}
 
     @staticmethod
-    def _get_parameter_scaler(
+    def _make_parameter_scaler(
         parameter: Parameter,
     ) -> ScalerProtocol | Literal["passthrough"]:
         # See base class.
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 6ef9e0196..07a14b66a 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -49,7 +49,7 @@ class RandomForestSurrogate(GaussianSurrogate):
     """The actual model."""
 
     @staticmethod
-    def _get_parameter_scaler(
+    def _make_parameter_scaler(
         parameter: Parameter,
     ) -> ScalerProtocol | Literal["passthrough"]:
         # See base class.

From 64b5450d5b6ec315ff846de1bc409a6b66dfe968 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 10:07:52 +0200
Subject: [PATCH 27/97] Draft output scaling mechanism

---
 baybe/surrogates/base.py | 70 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 67 insertions(+), 3 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 21e268763..5f802d038 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -4,6 +4,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
+from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 import pandas as pd
@@ -33,6 +34,7 @@
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
+    from botorch.models.transforms.outcome import OutcomeTransform
     from botorch.posteriors import GPyTorchPosterior, Posterior
     from sklearn.compose import ColumnTransformer
     from torch import Tensor
@@ -50,6 +52,16 @@
 """
 
 
+class _NoTransform(Enum):
+    """Sentinel class."""
+
+    IDENTITY_TRANSFORM = auto()
+
+
+_IDENTITY_TRANSFORM = _NoTransform.IDENTITY_TRANSFORM
+"""Sentinel to indicate the absence of a transform where `None` is ambiguous."""
+
+
 @define
 class Surrogate(ABC, SerialMixin):
     """Abstract base class for all surrogate models."""
@@ -80,6 +92,15 @@ class Surrogate(ABC, SerialMixin):
     representation to a corresponding dataframe containing their computational
     representation. Only available after the surrogate has been fitted."""
 
+    # TODO: type should be `Standardize | _NoTransform`` but is currently
+    #   omitted due to: https://github.com/python-attrs/cattrs/issues/531
+    _output_scaler = field(init=False, default=None, eq=False)
+    """Optional callable for scaling target values.
+
+    Scales a tensor containing target measurements in computational representation
+    to make them ready for processing by the surrogate. Only available after the
+    surrogate has been fitted."""
+
     def to_botorch(self) -> Model:
         """Create the botorch-ready representation of the model."""
         from baybe.surrogates._adapter import AdapterModel
@@ -93,10 +114,18 @@ def _make_parameter_scaler(
         """Return the scaler to be used for the given parameter."""
         return MinMaxScaler()
 
+    @staticmethod
+    def _make_target_scaler() -> OutcomeTransform | None:
+        """Return the scaler to be used for target scaling."""
+        from botorch.models.transforms.outcome import Standardize
+
+        # TODO: Multi-target extension
+        return Standardize(1)
+
     def _make_input_scaler(
         self, searchspace: SearchSpace, measurements: pd.DataFrame
     ) -> ColumnTransformer:
-        """Make a scaler to be used for transforming computational dataframes."""
+        """Make the input scaler for transforming computational dataframes."""
         from sklearn.compose import make_column_transformer
 
         # Create the composite scaler from the parameter-wise scaler objects
@@ -115,6 +144,23 @@ def _make_input_scaler(
 
         return scaler
 
+    def _make_output_scaler(
+        self, objective: Objective, measurements: pd.DataFrame
+    ) -> OutcomeTransform | _NoTransform:
+        """Make the output scaler for transforming computational dataframes."""
+        import torch
+
+        scaler = self._make_target_scaler()
+        if scaler is None:
+            return _IDENTITY_TRANSFORM
+
+        # TODO: Decide whether scaler is to be fit to target bounds and/or
+        #   extreme points in the given measurement data
+        scaler(torch.from_numpy(objective.transform(measurements).values))
+        scaler.eval()
+
+        return scaler
+
     def transform_inputs(self, data: pd.DataFrame) -> pd.DataFrame:
         """Transform an experimental parameter dataframe."""
         if self._input_transform is None:
@@ -129,7 +175,10 @@ def transform_targets(self, data: pd.DataFrame) -> pd.DataFrame:
 
     def posterior(self, candidates: pd.DataFrame) -> Posterior:
         """Evaluate the surrogate model at the given candidate points."""
-        return self._posterior(to_tensor(self.transform_inputs(candidates)))
+        p = self._posterior(to_tensor(self.transform_inputs(candidates)))
+        if self._output_scaler is not _IDENTITY_TRANSFORM:
+            p = self._output_scaler.untransform_posterior(p)
+        return p
 
     @abstractmethod
     def _posterior(self, candidates: Tensor) -> Posterior:
@@ -179,7 +228,10 @@ def fit(
                 "Continuous search spaces are currently only supported by GPs."
             )
 
+        # Create scaler objects
         input_scaler = self._make_input_scaler(searchspace, measurements)
+        output_scaler = self._make_output_scaler(objective, measurements)
+        self._output_scaler = output_scaler
 
         def transform_inputs(df: pd.DataFrame, /) -> pd.DataFrame:
             """Fitted input transformation pipeline."""
@@ -188,9 +240,21 @@ def transform_inputs(df: pd.DataFrame, /) -> pd.DataFrame:
                 out, index=df.index, columns=input_scaler.get_feature_names_out()
             )
 
+        def transform_outputs(df: pd.DataFrame, /) -> pd.DataFrame:
+            """Fitted output transformation pipeline."""
+            import torch
+
+            dft = objective.transform(df)
+
+            if self._output_scaler is _IDENTITY_TRANSFORM:
+                return dft
+
+            out = output_scaler(torch.from_numpy(dft.values))[0]
+            return pd.DataFrame(out.numpy(), index=df.index, columns=dft.columns)
+
         # Store context-specific transformations
         self._input_transform = transform_inputs
-        self._target_transform = lambda x: objective.transform(x)
+        self._target_transform = transform_outputs
 
         # Transform and fit
         train_x, train_y = to_tensor(

From 6dad04af2842771a8d08ed45e806f6403f257a11 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 10:10:34 +0200
Subject: [PATCH 28/97] Silence warning by allowing extra columns

---
 baybe/surrogates/base.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 5f802d038..084ad58c0 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -235,7 +235,9 @@ def fit(
 
         def transform_inputs(df: pd.DataFrame, /) -> pd.DataFrame:
             """Fitted input transformation pipeline."""
-            out = input_scaler.transform(searchspace.transform(df, allow_missing=True))
+            out = input_scaler.transform(
+                searchspace.transform(df, allow_extra=True, allow_missing=True)
+            )
             return pd.DataFrame(
                 out, index=df.index, columns=input_scaler.get_feature_names_out()
             )

From 25e356a91d2617b294459e93a34ad6b5f6bb4961 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 15:07:39 +0200
Subject: [PATCH 29/97] Improve signatures

---
 baybe/surrogates/base.py                         | 16 ++++++++--------
 baybe/surrogates/custom.py                       |  5 +++--
 baybe/surrogates/gaussian_process/core.py        |  2 +-
 baybe/surrogates/linear.py                       |  3 ++-
 baybe/surrogates/naive.py                        |  2 +-
 baybe/surrogates/ngboost.py                      |  2 +-
 baybe/surrogates/random_forest.py                |  2 +-
 baybe/surrogates/utils.py                        |  2 +-
 .../custom_architecture_sklearn.py               |  2 +-
 .../custom_architecture_torch.py                 |  2 +-
 10 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 084ad58c0..e7c687466 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -161,19 +161,19 @@ def _make_output_scaler(
 
         return scaler
 
-    def transform_inputs(self, data: pd.DataFrame) -> pd.DataFrame:
+    def transform_inputs(self, df: pd.DataFrame, /) -> pd.DataFrame:
         """Transform an experimental parameter dataframe."""
         if self._input_transform is None:
             raise ModelNotTrainedError("The model must be trained first.")
-        return self._input_transform(data)
+        return self._input_transform(df)
 
-    def transform_targets(self, data: pd.DataFrame) -> pd.DataFrame:
+    def transform_targets(self, df: pd.DataFrame, /) -> pd.DataFrame:
         """Transform an experimental measurement dataframe."""
         if self._target_transform is None:
             raise ModelNotTrainedError("The model must be trained first.")
-        return self._target_transform(data)
+        return self._target_transform(df)
 
-    def posterior(self, candidates: pd.DataFrame) -> Posterior:
+    def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
         """Evaluate the surrogate model at the given candidate points."""
         p = self._posterior(to_tensor(self.transform_inputs(candidates)))
         if self._output_scaler is not _IDENTITY_TRANSFORM:
@@ -181,7 +181,7 @@ def posterior(self, candidates: pd.DataFrame) -> Posterior:
         return p
 
     @abstractmethod
-    def _posterior(self, candidates: Tensor) -> Posterior:
+    def _posterior(self, candidates: Tensor, /) -> Posterior:
         """Perform the actual posterior evaluation logic."""
 
     @staticmethod
@@ -274,7 +274,7 @@ def _fit(self, train_x: Tensor, train_y: Tensor, context: Any = None) -> None:
 class GaussianSurrogate(Surrogate, ABC):
     """A surrogate model providing Gaussian posterior estimates."""
 
-    def _posterior(self, candidates: Tensor) -> GPyTorchPosterior:
+    def _posterior(self, candidates: Tensor, /) -> GPyTorchPosterior:
         # See base class.
 
         import torch
@@ -289,7 +289,7 @@ def _posterior(self, candidates: Tensor) -> GPyTorchPosterior:
         return GPyTorchPosterior(mvn)
 
     @abstractmethod
-    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         """Estimate first and second moments of the Gaussian posterior.
 
         The second moment may either be a 1-D tensor of marginal variances for the
diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py
index 9ed19e6c8..ba5291fe7 100644
--- a/baybe/surrogates/custom.py
+++ b/baybe/surrogates/custom.py
@@ -7,6 +7,7 @@
 
 It is planned to solve this issue in the future.
 """
+
 from __future__ import annotations
 
 from collections.abc import Callable
@@ -69,7 +70,7 @@ def __init__(self, *args, **kwargs):
             def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
                 return self._model._fit(train_x, train_y, context)
 
-            def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+            def _posterior(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
                 return self._model._posterior(candidates)
 
             def __get_attribute__(self, attr):
@@ -147,7 +148,7 @@ def default_model(self) -> ort.InferenceSession:
             raise ValueError("Invalid ONNX string") from exc
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         import torch
 
         from baybe.utils.torch import DTypeFloatTorch
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 1ff56c1d1..aec08b0be 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -126,7 +126,7 @@ def _get_model_context(
         # See base class.
         return _ModelContext(searchspace=searchspace)
 
-    def _posterior(self, candidates: Tensor) -> Posterior:
+    def _posterior(self, candidates: Tensor, /) -> Posterior:
         # See base class.
         return self._model.posterior(candidates)
 
diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py
index d523285b9..bc2f47ad1 100644
--- a/baybe/surrogates/linear.py
+++ b/baybe/surrogates/linear.py
@@ -6,6 +6,7 @@
 Since we plan to refactor the surrogates, this part of the documentation will be
 available in the future. Thus, please have a look in the source code directly.
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, ClassVar
@@ -45,7 +46,7 @@ class BayesianLinearSurrogate(GaussianSurrogate):
     """The actual model."""
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
diff --git a/baybe/surrogates/naive.py b/baybe/surrogates/naive.py
index 629ea5597..201275321 100644
--- a/baybe/surrogates/naive.py
+++ b/baybe/surrogates/naive.py
@@ -33,7 +33,7 @@ class MeanPredictionSurrogate(GaussianSurrogate):
     """The estimated posterior mean value of the training targets."""
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 6a4f98d2f..652cbfdb9 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -63,7 +63,7 @@ def _make_parameter_scaler(
         return "passthrough"
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 07a14b66a..0f45d6f79 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -58,7 +58,7 @@ def _make_parameter_scaler(
         return "passthrough"
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index 6958c36e4..89b1dd24f 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -45,7 +45,7 @@ def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
     _fit_original = cls._fit
     _posterior_original = cls._posterior
 
-    def _posterior_new(self, candidates: Tensor) -> Posterior:
+    def _posterior_new(self, candidates: Tensor, /) -> Posterior:
         """Use fallback model if it exists, otherwise call original posterior."""
         # Alternative model fallback
         if constant_target_model := _constant_target_model_store.get(id(self), None):
diff --git a/examples/Custom_Surrogates/custom_architecture_sklearn.py b/examples/Custom_Surrogates/custom_architecture_sklearn.py
index 0cf79741b..5a6b2e6f1 100644
--- a/examples/Custom_Surrogates/custom_architecture_sklearn.py
+++ b/examples/Custom_Surrogates/custom_architecture_sklearn.py
@@ -75,7 +75,7 @@ class StackingRegressorSurrogate:
     def __init__(self):
         self.model: StackingRegressor | None = None
 
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _posterior(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         """See :class:`baybe.surrogates.Surrogate`."""
         return self.model.predict(candidates)
 
diff --git a/examples/Custom_Surrogates/custom_architecture_torch.py b/examples/Custom_Surrogates/custom_architecture_torch.py
index b78eea581..6d04c9f87 100644
--- a/examples/Custom_Surrogates/custom_architecture_torch.py
+++ b/examples/Custom_Surrogates/custom_architecture_torch.py
@@ -115,7 +115,7 @@ class NeuralNetDropoutSurrogate:
     def __init__(self):
         self.model: nn.Module | None = None
 
-    def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]:
+    def _posterior(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         """See :class:`baybe.surrogates.Surrogate`."""
         self.model = self.model.train()  # keep dropout
         # Convert input from double to float

From 2a2849be2b49bd345579f92781083d311c880e50 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 15:31:25 +0200
Subject: [PATCH 30/97] Harmonize terminology

---
 baybe/acquisition/base.py |  2 +-
 baybe/surrogates/base.py  | 16 ++++++++--------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index cacbc77ef..ef67a4572 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -52,7 +52,7 @@ def to_botorch(
         params_dict = filter_attributes(object=self, callable_=acqf_cls.__init__)
 
         train_x = surrogate.transform_inputs(measurements)
-        train_y = surrogate.transform_targets(measurements)
+        train_y = surrogate.transform_outputs(measurements)
 
         signature_params = signature(acqf_cls).parameters
         additional_params = {}
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index e7c687466..f82c47058 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -83,10 +83,10 @@ class Surrogate(ABC, SerialMixin):
     representation to a corresponding dataframe containing their computational
     representation. Only available after the surrogate has been fitted."""
 
-    _target_transform: Callable[[pd.DataFrame], pd.DataFrame] | None = field(
+    _output_transform: Callable[[pd.DataFrame], pd.DataFrame] | None = field(
         init=False, default=None, eq=False
     )
-    """Callable preparing surrogate targets for training.
+    """Callable preparing surrogate outputs for training.
 
     Transforms a dataframe containing target measurements in experimental
     representation to a corresponding dataframe containing their computational
@@ -95,7 +95,7 @@ class Surrogate(ABC, SerialMixin):
     # TODO: type should be `Standardize | _NoTransform`` but is currently
     #   omitted due to: https://github.com/python-attrs/cattrs/issues/531
     _output_scaler = field(init=False, default=None, eq=False)
-    """Optional callable for scaling target values.
+    """Optional callable for scaling output values.
 
     Scales a tensor containing target measurements in computational representation
     to make them ready for processing by the surrogate. Only available after the
@@ -167,11 +167,11 @@ def transform_inputs(self, df: pd.DataFrame, /) -> pd.DataFrame:
             raise ModelNotTrainedError("The model must be trained first.")
         return self._input_transform(df)
 
-    def transform_targets(self, df: pd.DataFrame, /) -> pd.DataFrame:
+    def transform_outputs(self, df: pd.DataFrame, /) -> pd.DataFrame:
         """Transform an experimental measurement dataframe."""
-        if self._target_transform is None:
+        if self._output_transform is None:
             raise ModelNotTrainedError("The model must be trained first.")
-        return self._target_transform(df)
+        return self._output_transform(df)
 
     def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
         """Evaluate the surrogate model at the given candidate points."""
@@ -256,12 +256,12 @@ def transform_outputs(df: pd.DataFrame, /) -> pd.DataFrame:
 
         # Store context-specific transformations
         self._input_transform = transform_inputs
-        self._target_transform = transform_outputs
+        self._output_transform = transform_outputs
 
         # Transform and fit
         train_x, train_y = to_tensor(
             self.transform_inputs(measurements),
-            self.transform_targets(measurements),
+            self.transform_outputs(measurements),
         )
         self._fit(train_x, train_y, self._get_model_context(searchspace, objective))
 

From 920b079113a059426b54a6b988836041be771055 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 18:08:32 +0200
Subject: [PATCH 31/97] Update test for empty bounds

---
 tests/test_searchspace.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_searchspace.py b/tests/test_searchspace.py
index 81c083b14..23d2f8a9e 100644
--- a/tests/test_searchspace.py
+++ b/tests/test_searchspace.py
@@ -58,9 +58,9 @@ def test_empty_parameter_bounds():
     parameters = []
     searchspace_discrete = SubspaceDiscrete.from_product(parameters=parameters)
     searchspace_continuous = SubspaceContinuous(parameters=parameters)
-    expected = np.empty((2, 0))
-    assert np.array_equal(searchspace_discrete.comp_rep_bounds.values, expected)
-    assert np.array_equal(searchspace_continuous.comp_rep_bounds.values, expected)
+    expected = pd.DataFrame(np.empty((2, 0)), index=["min", "max"])
+    pd.testing.assert_frame_equal(searchspace_discrete.comp_rep_bounds, expected)
+    pd.testing.assert_frame_equal(searchspace_continuous.comp_rep_bounds, expected)
 
 
 def test_discrete_searchspace_creation_from_dataframe():

From cdf6688c2548754755a0a73bd2bebe25e50e213f Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 18:14:54 +0200
Subject: [PATCH 32/97] Fix import order

Became necessary due to the known-third-party ruff flag
---
 streamlit/initial_recommender.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/streamlit/initial_recommender.py b/streamlit/initial_recommender.py
index 44362c358..dcd77ac1f 100644
--- a/streamlit/initial_recommender.py
+++ b/streamlit/initial_recommender.py
@@ -3,9 +3,9 @@
 import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
+import streamlit as st
 from sklearn.datasets import make_blobs
 
-import streamlit as st
 from baybe.recommenders.pure.nonpredictive.base import NonPredictiveRecommender
 from baybe.searchspace import SearchSpace, SubspaceDiscrete
 from baybe.utils.basic import get_subclasses

From 6e052f71eb1bd2bd3fd7c8764fde85ad049932c3 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 21:44:22 +0200
Subject: [PATCH 33/97] Decide for transformation approach

For this PR, we leave the mechanism untouched:
* Parameters are normalized based on search space bounds
* Targets are standardized based on observed measurements
---
 baybe/surrogates/base.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index f82c47058..5fed29c65 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -122,9 +122,7 @@ def _make_target_scaler() -> OutcomeTransform | None:
         # TODO: Multi-target extension
         return Standardize(1)
 
-    def _make_input_scaler(
-        self, searchspace: SearchSpace, measurements: pd.DataFrame
-    ) -> ColumnTransformer:
+    def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
         """Make the input scaler for transforming computational dataframes."""
         from sklearn.compose import make_column_transformer
 
@@ -138,8 +136,6 @@ def _make_input_scaler(
         ]
         scaler = make_column_transformer(*transformers, verbose_feature_names_out=False)
 
-        # TODO: Decide whether scaler is to be fit to parameter bounds and/or
-        #   extreme points in the given measurement data
         scaler.fit(searchspace.comp_rep_bounds)
 
         return scaler
@@ -154,8 +150,7 @@ def _make_output_scaler(
         if scaler is None:
             return _IDENTITY_TRANSFORM
 
-        # TODO: Decide whether scaler is to be fit to target bounds and/or
-        #   extreme points in the given measurement data
+        # TODO: Consider taking into account target boundaries when available
         scaler(torch.from_numpy(objective.transform(measurements).values))
         scaler.eval()
 
@@ -229,7 +224,7 @@ def fit(
             )
 
         # Create scaler objects
-        input_scaler = self._make_input_scaler(searchspace, measurements)
+        input_scaler = self._make_input_scaler(searchspace)
         output_scaler = self._make_output_scaler(objective, measurements)
         self._output_scaler = output_scaler
 

From ef84a35cea0370f1140bdfa3758119b42eda142e Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 22 Jul 2024 21:49:47 +0200
Subject: [PATCH 34/97] Update docstrings

---
 baybe/surrogates/base.py | 8 ++++----
 baybe/utils/scaling.py   | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 5fed29c65..0394ae51e 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -80,8 +80,8 @@ class Surrogate(ABC, SerialMixin):
     """Callable preparing surrogate inputs for training/prediction.
 
     Transforms a dataframe containing parameter configurations in experimental
-    representation to a corresponding dataframe containing their computational
-    representation. Only available after the surrogate has been fitted."""
+    representation to a corresponding dataframe containing their **scaled**
+    computational representation. Only available after the surrogate has been fitted."""
 
     _output_transform: Callable[[pd.DataFrame], pd.DataFrame] | None = field(
         init=False, default=None, eq=False
@@ -89,8 +89,8 @@ class Surrogate(ABC, SerialMixin):
     """Callable preparing surrogate outputs for training.
 
     Transforms a dataframe containing target measurements in experimental
-    representation to a corresponding dataframe containing their computational
-    representation. Only available after the surrogate has been fitted."""
+    representation to a corresponding dataframe containing their **scaled**
+    computational representation. Only available after the surrogate has been fitted."""
 
     # TODO: type should be `Standardize | _NoTransform`` but is currently
     #   omitted due to: https://github.com/python-attrs/cattrs/issues/531
diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index 0bff2f941..c84a6adca 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -16,7 +16,7 @@ class ScalerProtocol(Protocol):
     """
 
     def fit(self, df: pd.DataFrame, /) -> None:
-        """Fit the scaler to a given data set."""
+        """Fit the scaler to a given dataframe."""
 
     def transform(self, df: pd.DataFrame, /) -> pd.DataFrame:
-        """Transform a data using the fitted scaling logic."""
+        """Transform a dataframe using the fitted scaling logic."""

From 2b3dcabc875e14175739b133fa3cff27a043b212 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 23 Jul 2024 09:07:36 +0200
Subject: [PATCH 35/97] Remove separate scaling logic from GPs

---
 baybe/surrogates/gaussian_process/core.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index aec08b0be..6480ed7fe 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -6,9 +6,11 @@
 
 from attrs import define, field
 from attrs.validators import instance_of
+from sklearn.preprocessing import MinMaxScaler
 
 from baybe.objective import Objective
 from baybe.parameters.base import Parameter
+from baybe.parameters.categorical import TaskParameter
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.base import Surrogate
 from baybe.surrogates.gaussian_process.kernel_factory import (
@@ -116,8 +118,11 @@ def _make_parameter_scaler(
     ) -> ScalerProtocol | Literal["passthrough"]:
         # See base class.
 
-        # For GPs, we use botorch's built-in machinery for scaling
-        return "passthrough"
+        # Task parameters are handled separately through an index kernel
+        if isinstance(parameter, TaskParameter):
+            return "passthrough"
+
+        return MinMaxScaler()
 
     @staticmethod
     def _get_model_context(
@@ -139,14 +144,6 @@ def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None
 
         numerical_idxs = context.get_numerical_indices(train_x.shape[-1])
 
-        # define the input and outcome transforms
-        # TODO [Scaling]: scaling should be handled by search space object
-        input_transform = botorch.models.transforms.Normalize(
-            train_x.shape[1], bounds=context.parameter_bounds, indices=numerical_idxs
-        )
-        # TODO: use target value bounds when explicitly provided
-        outcome_transform = botorch.models.transforms.Standardize(train_y.shape[1])
-
         # extract the batch shape of the training data
         batch_shape = train_x.shape[:-2]
 
@@ -184,8 +181,6 @@ def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None
         self._model = botorch.models.SingleTaskGP(
             train_x,
             train_y,
-            input_transform=input_transform,
-            outcome_transform=outcome_transform,
             mean_module=mean_module,
             covar_module=covar_module,
             likelihood=likelihood,

From 161bddbdaa95d8fed65b5a61acbe92e65cc73480 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 23 Jul 2024 09:54:35 +0200
Subject: [PATCH 36/97] Rename ScalerProtocol to ParameterScalerProtocol

In order to differentiate from target scaling
---
 baybe/surrogates/base.py                  | 4 ++--
 baybe/surrogates/gaussian_process/core.py | 4 ++--
 baybe/surrogates/ngboost.py               | 4 ++--
 baybe/surrogates/random_forest.py         | 4 ++--
 baybe/utils/scaling.py                    | 8 ++++----
 5 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 0394ae51e..346ae3e24 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -30,7 +30,7 @@
 )
 from baybe.serialization.mixin import SerialMixin
 from baybe.utils.dataframe import to_tensor
-from baybe.utils.scaling import ScalerProtocol
+from baybe.utils.scaling import ParameterScalerProtocol
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
@@ -110,7 +110,7 @@ def to_botorch(self) -> Model:
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | Literal["passthrough"]:
         """Return the scaler to be used for the given parameter."""
         return MinMaxScaler()
 
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 6480ed7fe..f727b96e9 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -25,7 +25,7 @@
     DefaultKernelFactory,
     _default_noise_factory,
 )
-from baybe.utils.scaling import ScalerProtocol
+from baybe.utils.scaling import ParameterScalerProtocol
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
@@ -115,7 +115,7 @@ def to_botorch(self) -> Model:  # noqa: D102
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | Literal["passthrough"]:
         # See base class.
 
         # Task parameters are handled separately through an index kernel
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 652cbfdb9..e0072da54 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -18,7 +18,7 @@
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
-from baybe.utils.scaling import ScalerProtocol
+from baybe.utils.scaling import ParameterScalerProtocol
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -56,7 +56,7 @@ def __attrs_post_init__(self):
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | Literal["passthrough"]:
         # See base class.
 
         # Tree-like models do not require any input scaling
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 0f45d6f79..d7b1df9c2 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -19,7 +19,7 @@
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
-from baybe.utils.scaling import ScalerProtocol
+from baybe.utils.scaling import ParameterScalerProtocol
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -51,7 +51,7 @@ class RandomForestSurrogate(GaussianSurrogate):
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | Literal["passthrough"]:
         # See base class.
 
         # Tree-like models do not require any input scaling
diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index c84a6adca..3c1b419cf 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -7,8 +7,8 @@
 import pandas as pd
 
 
-class ScalerProtocol(Protocol):
-    """Type protocol specifying the interface scalers need to implement.
+class ParameterScalerProtocol(Protocol):
+    """Type protocol specifying the interface parameter scalers need to implement.
 
     The protocol is compatible with sklearn scalers such as
     :class:`sklearn.preprocessing.MinMaxScaler` or
@@ -16,7 +16,7 @@ class ScalerProtocol(Protocol):
     """
 
     def fit(self, df: pd.DataFrame, /) -> None:
-        """Fit the scaler to a given dataframe."""
+        """Fit the scaler to a given dataframe containing parameter configurations."""
 
     def transform(self, df: pd.DataFrame, /) -> pd.DataFrame:
-        """Transform a dataframe using the fitted scaling logic."""
+        """Transform a parameter dataframe using the fitted scaling logic."""

From e7f3f67ec114fdf500c567c162e43b761ed7540a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 23 Jul 2024 09:51:31 +0200
Subject: [PATCH 37/97] Update CHANGELOG.md

---
 CHANGELOG.md | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dc4bf2504..6d1966232 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -37,6 +37,8 @@ _ `_optional` subpackage for managing optional dependencies
 - `GaussianSurrogate` base class for surrogate models with Gaussian posteriors
 - `comp_rep_columns` property for `Parameter`, `SearchSpace`, `SubspaceDiscrete`
   and `SubspaceContinuous` classes
+- Reworked mechanisms for surrogate input/output scaling configurable per class
+- `ParameterScalerProtocol` class for enabling user-defined input scaling mechanisms
 
 ### Changed
 - Passing an `Objective` to `Campaign` is now optional
@@ -48,11 +50,12 @@ _ `_optional` subpackage for managing optional dependencies
 - Context information required by `Surrogate` models is now cleanly encapsulated into
   a `context` object passed to `Surrogate._fit`
 - Fallback models created by `catch_constant_targets` are stored outside of surrogate
+- `to_tensor` now also handles `numpy` arrays
+- `GaussianProcessSurrogate` no longer uses a separate scaling approach
 
 ### Removed
 - Support for Python 3.9 removed due to new [BoTorch requirements](https://github.com/pytorch/botorch/pull/2293) 
   and guidelines from [Scientific Python](https://scientific-python.org/specs/spec-0000/)
-- `register_custom_architecture` decorator
 - `Scalar` and `DefaultScaler` classes
 
 ### Fixed
@@ -68,8 +71,6 @@ _ `_optional` subpackage for managing optional dependencies
 - Passing a dataframe via the `data` argument to the `transform` methods of
   `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous` is no longer possible.
   The dataframe must now be passed as positional argument.
-- Role of `register_custom_architecture` has been taken over by
-  `baybe.surrogates.base.SurrogateProtocol`
 
 ## [0.9.1] - 2024-06-04
 ### Changed

From 21953d44f5aa9b721fc20e246dccd62047d3012b Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 23 Jul 2024 13:14:33 +0200
Subject: [PATCH 38/97] Replace literal return type with None

---
 baybe/surrogates/base.py                  | 6 +++---
 baybe/surrogates/gaussian_process/core.py | 6 +++---
 baybe/surrogates/ngboost.py               | 6 +++---
 baybe/surrogates/random_forest.py         | 6 +++---
 4 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 346ae3e24..efc83e8da 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import pandas as pd
 from attrs import define, field
@@ -110,7 +110,7 @@ def to_botorch(self) -> Model:
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         """Return the scaler to be used for the given parameter."""
         return MinMaxScaler()
 
@@ -129,7 +129,7 @@ def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
         # Create the composite scaler from the parameter-wise scaler objects
         transformers = [
             (
-                self._make_parameter_scaler(p),
+                "passthrough" if (s := self._make_parameter_scaler(p)) is None else s,
                 [c for c in p.comp_rep_columns if c in searchspace.comp_rep_columns],
             )
             for p in searchspace.parameters
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index f727b96e9..5b00982e6 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar, Literal
+from typing import TYPE_CHECKING, ClassVar
 
 from attrs import define, field
 from attrs.validators import instance_of
@@ -115,12 +115,12 @@ def to_botorch(self) -> Model:  # noqa: D102
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         # See base class.
 
         # Task parameters are handled separately through an index kernel
         if isinstance(parameter, TaskParameter):
-            return "passthrough"
+            return
 
         return MinMaxScaler()
 
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index e0072da54..b0ce3849c 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from attr import define, field
 from ngboost import NGBRegressor
@@ -56,11 +56,11 @@ def __attrs_post_init__(self):
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return "passthrough"
+        return
 
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index d7b1df9c2..776a9a830 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 from attr import define, field
@@ -51,11 +51,11 @@ class RandomForestSurrogate(GaussianSurrogate):
     @staticmethod
     def _make_parameter_scaler(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | Literal["passthrough"]:
+    ) -> ParameterScalerProtocol | None:
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return "passthrough"
+        return
 
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:

From 536a3a881f73896a6b0d98879c86f525d783b91a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 24 Jul 2024 10:52:05 +0200
Subject: [PATCH 39/97] Implement workaround to circumvent ColumnTransformer
 limitations

---
 baybe/surrogates/base.py | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index efc83e8da..582fe0a2b 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -230,13 +230,42 @@ def fit(
 
         def transform_inputs(df: pd.DataFrame, /) -> pd.DataFrame:
             """Fitted input transformation pipeline."""
+            # IMPROVE: This method currently relies on two workarounds required
+            #   due the working mechanism of sklearn's ColumnTransformer:
+            #   * Unfortunately, the transformer returns a raw array, meaning that
+            #       column names need to be manually attached afterward.
+            #   * In certain cases (e.g. in hybrid spaces), the method needs
+            #       to transform only a subset of columns. Unfortunately, it is not
+            #       possible to use a subset of columns once the transformer is set up,
+            #       which is a side-effect of the first point. As a workaround,
+            #       we thus fill the missing columns with NaN and subselect afterward.
+
+            # For the workaround, collect all comp rep columns of the parameters
+            # that are actually present in the given dataframe. At the end,
+            # we'll filter the transformed augmented dataframe down to these columns.
+            exp_rep_cols = [p.name for p in searchspace.parameters]
+            comp_rep_cols = []
+            for col in [c for c in df.columns if c in exp_rep_cols]:
+                parameter = next(p for p in searchspace.parameters if p.name == col)
+                comp_rep_cols.extend(parameter.comp_rep_columns)
+
+            # Actual workaround: augment the dataframe with NaN for missing parameters
+            df_augmented = df.reindex(columns=exp_rep_cols)
+
+            # The actual transformation step
             out = input_scaler.transform(
-                searchspace.transform(df, allow_extra=True, allow_missing=True)
+                searchspace.transform(df_augmented, allow_extra=True)
             )
-            return pd.DataFrame(
+            out = pd.DataFrame(
                 out, index=df.index, columns=input_scaler.get_feature_names_out()
             )
 
+            # Undo the augmentation, taking into account that not all comp rep
+            # parameter columns may actually be part of the search space due
+            # to other preprocessing steps.
+            comp_rep_cols = list(set(comp_rep_cols).intersection(out.columns))
+            return out[comp_rep_cols]
+
         def transform_outputs(df: pd.DataFrame, /) -> pd.DataFrame:
             """Fitted output transformation pipeline."""
             import torch

From b88b3ba56045b31adf4f2b1e3d7a24c71d5d16e6 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 24 Jul 2024 10:52:53 +0200
Subject: [PATCH 40/97] Improve code grouping

---
 baybe/surrogates/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 582fe0a2b..27622d121 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -226,7 +226,6 @@ def fit(
         # Create scaler objects
         input_scaler = self._make_input_scaler(searchspace)
         output_scaler = self._make_output_scaler(objective, measurements)
-        self._output_scaler = output_scaler
 
         def transform_inputs(df: pd.DataFrame, /) -> pd.DataFrame:
             """Fitted input transformation pipeline."""
@@ -272,7 +271,7 @@ def transform_outputs(df: pd.DataFrame, /) -> pd.DataFrame:
 
             dft = objective.transform(df)
 
-            if self._output_scaler is _IDENTITY_TRANSFORM:
+            if output_scaler is _IDENTITY_TRANSFORM:
                 return dft
 
             out = output_scaler(torch.from_numpy(dft.values))[0]
@@ -281,6 +280,7 @@ def transform_outputs(df: pd.DataFrame, /) -> pd.DataFrame:
         # Store context-specific transformations
         self._input_transform = transform_inputs
         self._output_transform = transform_outputs
+        self._output_scaler = output_scaler
 
         # Transform and fit
         train_x, train_y = to_tensor(

From 1619bd78c46c6c11d4c7b3a43a118e0b12d99732 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 24 Jul 2024 13:48:18 +0200
Subject: [PATCH 41/97] Remove register_custom_architecture decorator

The decorator is no longer compatible with the generalized surrogate
layout. Instead of upgrading the decorator, a replacement mechanism
using Python's built-in typing.Protocol will be introduced.
---
 CHANGELOG.md                                  |   2 +
 baybe/surrogates/custom.py                    |  90 +------
 .../Custom_Surrogates_Header.md               |   2 +-
 .../custom_architecture_sklearn.py            | 171 -------------
 .../custom_architecture_torch.py              | 227 ------------------
 tests/test_custom_surrogate.py                |  47 +---
 tests/test_deprecations.py                    |   8 +
 7 files changed, 22 insertions(+), 525 deletions(-)
 delete mode 100644 examples/Custom_Surrogates/custom_architecture_sklearn.py
 delete mode 100644 examples/Custom_Surrogates/custom_architecture_torch.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6d1966232..d379feb50 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -56,6 +56,7 @@ _ `_optional` subpackage for managing optional dependencies
 ### Removed
 - Support for Python 3.9 removed due to new [BoTorch requirements](https://github.com/pytorch/botorch/pull/2293) 
   and guidelines from [Scientific Python](https://scientific-python.org/specs/spec-0000/)
+- `register_custom_architecture` decorator
 - `Scalar` and `DefaultScaler` classes
 
 ### Fixed
@@ -71,6 +72,7 @@ _ `_optional` subpackage for managing optional dependencies
 - Passing a dataframe via the `data` argument to the `transform` methods of
   `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous` is no longer possible.
   The dataframe must now be passed as positional argument.
+- A deprecation error is thrown when attempting to use `register_custom_architecture`
 
 ## [0.9.1] - 2024-06-04
 ### Changed
diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py
index ba5291fe7..2b2086c75 100644
--- a/baybe/surrogates/custom.py
+++ b/baybe/surrogates/custom.py
@@ -10,11 +10,11 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, ClassVar, NoReturn
 
 from attrs import define, field, validators
 
+from baybe.exceptions import DeprecationError
 from baybe.parameters import (
     CategoricalEncoding,
     CategoricalParameter,
@@ -24,10 +24,8 @@
     TaskParameter,
 )
 from baybe.searchspace import SearchSpace
-from baybe.serialization.core import block_serialization_hook, converter
-from baybe.surrogates.base import GaussianSurrogate, Surrogate
-from baybe.surrogates.utils import batchify, catch_constant_targets
-from baybe.surrogates.validation import validate_custom_architecture_cls
+from baybe.surrogates.base import GaussianSurrogate
+from baybe.surrogates.utils import batchify
 from baybe.utils.numerical import DTypeFloatONNX
 
 if TYPE_CHECKING:
@@ -35,80 +33,12 @@
     from torch import Tensor
 
 
-def register_custom_architecture(
-    joint_posterior_attr: bool = False,
-    constant_target_catching: bool = True,
-    batchify_posterior: bool = True,
-) -> Callable:
-    """Wrap a given custom model architecture class into a ```Surrogate```.
-
-    Args:
-        joint_posterior_attr: Boolean indicating if the model returns a posterior
-            distribution jointly across candidates or on individual points.
-        constant_target_catching: Boolean indicating if the model cannot handle
-            constant target values and needs the @catch_constant_targets decorator.
-        batchify_posterior: Boolean indicating if the model is incompatible
-            with t- and q-batching and needs the @batchify decorator for its posterior.
-
-    Returns:
-        A function that wraps around a model class based on the specifications.
-    """
-
-    def construct_custom_architecture(model_cls):
-        """Construct a surrogate class wrapped around the custom class."""
-        validate_custom_architecture_cls(model_cls)
-
-        class CustomArchitectureSurrogate(Surrogate):
-            """Wraps around a custom architecture class."""
-
-            joint_posterior: ClassVar[bool] = joint_posterior_attr
-            supports_transfer_learning: ClassVar[bool] = False
-
-            def __init__(self, *args, **kwargs):
-                self._model = model_cls(*args, **kwargs)
-
-            def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
-                return self._model._fit(train_x, train_y, context)
-
-            def _posterior(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
-                return self._model._posterior(candidates)
-
-            def __get_attribute__(self, attr):
-                """Access the attributes of the class instance if available.
-
-                If the attributes are not available,
-                it uses the attributes of the internal model instance.
-                """
-                # Try to retrieve the attribute in the class
-                try:
-                    val = super().__getattribute__(attr)
-                except AttributeError:
-                    pass
-                else:
-                    return val
-
-                # If the attribute is not overwritten, use that of the internal model
-                return self._model.__getattribute__(attr)
-
-        # Catch constant targets if needed
-        cls = (
-            catch_constant_targets(CustomArchitectureSurrogate)
-            if constant_target_catching
-            else CustomArchitectureSurrogate
-        )
-
-        # batchify posterior if needed
-        if batchify_posterior:
-            cls._posterior = batchify(cls._posterior)
-
-        # Block serialization of custom architectures
-        converter.register_unstructure_hook(
-            CustomArchitectureSurrogate, block_serialization_hook
-        )
-
-        return cls
-
-    return construct_custom_architecture
+def register_custom_architecture(*args, **kwargs) -> NoReturn:
+    """Deprecated! Raises an error when used."""  # noqa: D401
+    raise DeprecationError(
+        "The 'register_custom_architecture' decorator is no longer available. "
+        "There will be a replacement based on Python's `typing.Protocol` soon."
+    )
 
 
 @define(kw_only=True)
diff --git a/examples/Custom_Surrogates/Custom_Surrogates_Header.md b/examples/Custom_Surrogates/Custom_Surrogates_Header.md
index 0ff68a217..b6297f5f0 100644
--- a/examples/Custom_Surrogates/Custom_Surrogates_Header.md
+++ b/examples/Custom_Surrogates/Custom_Surrogates_Header.md
@@ -1,3 +1,3 @@
 # Custom Surrogates
 
-These examples demonstrate how to use custom {doc}`Surrogate </userguide/surrogates>` architectures and pre-trained [ONNX](https://onnx.ai) models.
\ No newline at end of file
+These examples demonstrate how to use custom pre-trained [ONNX](https://onnx.ai) surrogate models.
\ No newline at end of file
diff --git a/examples/Custom_Surrogates/custom_architecture_sklearn.py b/examples/Custom_Surrogates/custom_architecture_sklearn.py
deleted file mode 100644
index 5a6b2e6f1..000000000
--- a/examples/Custom_Surrogates/custom_architecture_sklearn.py
+++ /dev/null
@@ -1,171 +0,0 @@
-## Example for surrogate model with a custom architecture using `sklearn`
-
-# This example shows how to define a `sklearn` model architecture and use it as a surrogate.
-# Please note that the model is not designed to be useful but to demonstrate the workflow.
-
-# This example assumes some basic familiarity with using BayBE.
-# We thus refer to [`campaign`](./../Basics/campaign.md) for a basic example.
-
-### Necessary imports
-
-
-import numpy as np
-import torch
-from sklearn.base import BaseEstimator, RegressorMixin
-from sklearn.ensemble import (
-    GradientBoostingRegressor,
-    RandomForestRegressor,
-    StackingRegressor,
-)
-from sklearn.linear_model import LinearRegression, Ridge
-from torch import Tensor
-
-from baybe.campaign import Campaign
-from baybe.objectives import SingleTargetObjective
-from baybe.parameters import (
-    CategoricalParameter,
-    NumericalDiscreteParameter,
-    SubstanceParameter,
-)
-from baybe.recommenders import (
-    BotorchRecommender,
-    FPSRecommender,
-    TwoPhaseMetaRecommender,
-)
-from baybe.searchspace import SearchSpace
-from baybe.surrogates import register_custom_architecture
-from baybe.targets import NumericalTarget
-from baybe.utils.dataframe import add_fake_results
-
-### Surrogate Definition with BayBE Registration
-
-# The final estimator class must follow the sklearn estimator interface.
-# More details [here](https://scikit-learn.org/stable/developers/develop.html).
-
-# The choice of using tensors in fit/predict is purely for BayBE, not a requirement.
-
-# Final estimator
-
-
-class MeanVarEstimator(BaseEstimator, RegressorMixin):
-    """Stack final estimator for mean and variance."""
-
-    def fit(self, data: Tensor, targets: Tensor) -> None:
-        """No fit needed."""
-        return
-
-    def predict(self, data: Tensor) -> tuple[Tensor, Tensor]:
-        """Predict based on ensemble unweighted mean and variance."""
-        mean = torch.tensor(data.mean(axis=1))
-        var = torch.tensor(data.var(axis=1))
-        return mean, var
-
-
-# Registration
-
-# The class must include `_fit` and `_posterior` functions with the correct signatures.
-
-
-@register_custom_architecture(
-    joint_posterior_attr=False, constant_target_catching=False, batchify_posterior=True
-)
-class StackingRegressorSurrogate:
-    """Surrogate that extracts posterior from a stack of different regressors."""
-
-    def __init__(self):
-        self.model: StackingRegressor | None = None
-
-    def _posterior(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
-        """See :class:`baybe.surrogates.Surrogate`."""
-        return self.model.predict(candidates)
-
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
-        """See :class:`baybe.surrogates.Surrogate`."""
-        estimators = [
-            ("rf", RandomForestRegressor()),
-            ("gb", GradientBoostingRegressor()),
-            ("lr", LinearRegression()),
-            ("rr", Ridge()),
-        ]
-
-        self.model = StackingRegressor(
-            estimators=estimators,
-            final_estimator=MeanVarEstimator(),
-            cv=2,
-        )
-
-        self.model.fit(train_x, train_y.ravel())
-
-
-### Experiment Setup
-
-parameters = [
-    CategoricalParameter(
-        name="Granularity",
-        values=["coarse", "medium", "fine"],
-        encoding="OHE",
-    ),
-    NumericalDiscreteParameter(
-        name="Pressure[bar]",
-        values=[1, 5, 10],
-        tolerance=0.2,
-    ),
-    NumericalDiscreteParameter(
-        name="Temperature[degree_C]",
-        values=np.linspace(100, 200, 10),
-    ),
-    SubstanceParameter(
-        name="Solvent",
-        data={
-            "Solvent A": "COC",
-            "Solvent B": "CCC",
-            "Solvent C": "O",
-            "Solvent D": "CS(=O)C",
-        },
-        encoding="MORDRED",
-    ),
-]
-
-
-### Run DOE iterations with custom surrogate
-# Create campaign
-
-campaign = Campaign(
-    searchspace=SearchSpace.from_product(parameters=parameters, constraints=None),
-    objective=SingleTargetObjective(target=NumericalTarget(name="Yield", mode="MAX")),
-    recommender=TwoPhaseMetaRecommender(
-        recommender=BotorchRecommender(surrogate_model=StackingRegressorSurrogate()),
-        initial_recommender=FPSRecommender(),
-    ),
-)
-
-# Let's do a first round of recommendation
-recommendation = campaign.recommend(batch_size=2)
-
-print("Recommendation from campaign:")
-print(recommendation)
-
-# Add some fake results
-
-add_fake_results(recommendation, campaign)
-campaign.add_measurements(recommendation)
-
-# Do another round of recommendations
-recommendation = campaign.recommend(batch_size=2)
-
-# Print second round of recommendations
-
-print("Recommendation from campaign:")
-print(recommendation)
-
-print()
-
-
-### Serialization
-
-# Serialization of custom models is not supported
-
-try:
-    campaign.to_json()
-except RuntimeError as e:
-    print(f"Serialization Error Message: {e}")
diff --git a/examples/Custom_Surrogates/custom_architecture_torch.py b/examples/Custom_Surrogates/custom_architecture_torch.py
deleted file mode 100644
index 6d04c9f87..000000000
--- a/examples/Custom_Surrogates/custom_architecture_torch.py
+++ /dev/null
@@ -1,227 +0,0 @@
-## Example for surrogate model with a custom architecture using `pytorch`
-
-# This example shows how to define a `pytorch` model architecture and use it as a surrogate.
-# Please note that the model is not designed to be useful but to demonstrate the workflow.
-
-# This example assumes some basic familiarity with using BayBE.
-# We thus refer to [`campaign`](./../Basics/campaign.md) for a basic example.
-
-### Necessary imports
-
-
-import numpy as np
-import torch
-from torch import Tensor, nn
-
-from baybe.campaign import Campaign
-from baybe.objectives import SingleTargetObjective
-from baybe.parameters import (
-    CategoricalParameter,
-    NumericalDiscreteParameter,
-    SubstanceParameter,
-)
-from baybe.recommenders import (
-    BotorchRecommender,
-    FPSRecommender,
-    TwoPhaseMetaRecommender,
-)
-from baybe.searchspace import SearchSpace
-from baybe.surrogates import register_custom_architecture
-from baybe.targets import NumericalTarget
-from baybe.utils.dataframe import add_fake_results
-
-### Architecture definition
-
-# Note that the following is an example `PyTorch` Neural Network.
-# Details of the setup is not the focus of BayBE but can be found in `Pytorch` guides.
-
-# Model Configuration
-
-INPUT_DIM = 10
-OUTPUT_DIM = 1
-DROPOUT = 0.5
-NUM_NEURONS = [128, 32, 8]
-
-# Model training hyperparameters
-
-HYPERPARAMS = {
-    "epochs": 10,
-    "lr": 1e-3,
-    "criterion": nn.MSELoss,
-    "optimizer": torch.optim.Adam,
-}
-
-# MC Parameters
-
-MC = 100
-
-
-# Helper functions
-
-
-def _create_linear_block(in_features: int, out_features: int) -> list:
-    """Create a linear block with dropout and relu activation."""
-    return [nn.Linear(in_features, out_features), nn.Dropout(p=DROPOUT), nn.ReLU()]
-
-
-def _create_hidden_layers(num_neurons: list[int]) -> list:
-    """Create all hidden layers comprised of linear blocks."""
-    layers = []
-    for in_features, out_features in zip(num_neurons, num_neurons[1:]):
-        layers.extend(_create_linear_block(in_features, out_features))
-
-    return layers
-
-
-# Model Architecture
-
-
-class NeuralNetDropout(nn.Module):
-    """Pytorch implementation of Neural Network with Dropout."""
-
-    def __init__(self):
-        super().__init__()
-        layers = [
-            # Initial linear block with input
-            *(_create_linear_block(INPUT_DIM, NUM_NEURONS[0])),
-            # All hidden layers
-            *(_create_hidden_layers(NUM_NEURONS)),
-            # Last linear output
-            nn.Linear(NUM_NEURONS[-1], OUTPUT_DIM),
-        ]
-
-        # Sequential with layers (Feed Forward)
-        self.model = nn.Sequential(*layers)
-
-    def forward(self, data: Tensor) -> Tensor:
-        """Forward method for NN."""
-        return self.model(data)
-
-
-### Surrogate Definition with BayBE Registration
-
-# The class must include `_fit` and `_posterior` functions with the correct signatures
-
-
-# Registration
-
-
-@register_custom_architecture(
-    joint_posterior_attr=False, constant_target_catching=False, batchify_posterior=True
-)
-class NeuralNetDropoutSurrogate:
-    """Surrogate that extracts posterior using monte carlo dropout simulations."""
-
-    def __init__(self):
-        self.model: nn.Module | None = None
-
-    def _posterior(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
-        """See :class:`baybe.surrogates.Surrogate`."""
-        self.model = self.model.train()  # keep dropout
-        # Convert input from double to float
-        candidates = candidates.float()
-        # Run mc experiments through the NN with dropout
-        predictions = torch.cat(
-            [self.model(candidates).unsqueeze(dim=0) for _ in range(MC)]
-        )
-
-        # Compute posterior mean and variance
-        mean = predictions.mean(dim=0)
-        var = predictions.var(dim=0)
-
-        return mean, var
-
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
-        """See :class:`baybe.surrogates.Surrogate`."""
-        # Initialize Model
-        self.model = NeuralNetDropout()
-
-        # Training hyperparameters
-        opt = HYPERPARAMS["optimizer"](self.model.parameters(), lr=HYPERPARAMS["lr"])
-        criterion = HYPERPARAMS["criterion"]()
-
-        # Convert input from double to float
-        train_x = train_x.float()
-        train_y = train_y.float()
-
-        # Training loop
-        for _ in range(HYPERPARAMS["epochs"]):
-            opt.zero_grad()
-            preds = self.model(train_x)
-            loss = criterion(preds, train_y)
-            loss.backward()
-            opt.step()
-
-
-### Experiment Setup
-
-parameters = [
-    CategoricalParameter(
-        name="Granularity",
-        values=["coarse", "medium", "fine"],
-        encoding="OHE",
-    ),
-    NumericalDiscreteParameter(
-        name="Pressure[bar]",
-        values=[1, 5, 10],
-        tolerance=0.2,
-    ),
-    NumericalDiscreteParameter(
-        name="Temperature[degree_C]",
-        values=np.linspace(100, 200, 10),
-    ),
-    SubstanceParameter(
-        name="Solvent",
-        data={
-            "Solvent A": "COC",
-            "Solvent B": "CCC",
-            "Solvent C": "O",
-            "Solvent D": "CS(=O)C",
-        },
-        encoding="MORDRED",
-    ),
-]
-
-
-### Run DOE iterations with custom surrogate
-# Create campaign
-
-campaign = Campaign(
-    searchspace=SearchSpace.from_product(parameters=parameters, constraints=None),
-    objective=SingleTargetObjective(target=NumericalTarget(name="Yield", mode="MAX")),
-    recommender=TwoPhaseMetaRecommender(
-        recommender=BotorchRecommender(surrogate_model=NeuralNetDropoutSurrogate()),
-        initial_recommender=FPSRecommender(),
-    ),
-)
-
-# Let's do a first round of recommendation
-recommendation = campaign.recommend(batch_size=2)
-
-print("Recommendation from campaign:")
-print(recommendation)
-
-# Add some fake results
-
-add_fake_results(recommendation, campaign)
-campaign.add_measurements(recommendation)
-
-# Do another round of recommendations
-recommendation = campaign.recommend(batch_size=2)
-
-# Print second round of recommendations
-
-print("Recommendation from campaign:")
-print(recommendation)
-
-print()
-
-
-### Serialization
-
-# Serialization of custom models is not supported
-
-try:
-    campaign.to_json()
-except RuntimeError as e:
-    print(f"Serialization Error Message: {e}")
diff --git a/tests/test_custom_surrogate.py b/tests/test_custom_surrogate.py
index e6411da0d..b34901914 100644
--- a/tests/test_custom_surrogate.py
+++ b/tests/test_custom_surrogate.py
@@ -5,7 +5,7 @@
 
 from baybe import Campaign
 from baybe._optional.info import ONNX_INSTALLED
-from baybe.surrogates import CustomONNXSurrogate, register_custom_architecture
+from baybe.surrogates import CustomONNXSurrogate
 from tests.conftest import run_iterations
 
 
@@ -50,48 +50,3 @@ def test_supported_parameter_types(campaign: Campaign, should_raise: bool):
     context = pytest.raises(TypeError) if should_raise else nullcontext()
     with context:
         campaign.recommend(batch_size=1)
-
-
-def test_validate_architectures():
-    """Test architecture class validation."""
-    # Scenario: Empty Class
-    with pytest.raises(ValueError):
-        register_custom_architecture()(str)
-
-    # Scenario: Class with just `_fit`
-    with pytest.raises(ValueError):
-        register_custom_architecture()(type("PartialArch", (), {"_fit": True}))
-
-    # Scenario: Class with `_fit` and `_posterior` but not methods
-    with pytest.raises(ValueError):
-        register_custom_architecture()(
-            type("PartialArch", (), {"_fit": True, "_posterior": True})
-        )
-
-    # Scenario: Class with invalid `_fit` and `_posterior` methods
-    def _invalid_func(invalid_param1, invalid_param2=1):
-        return invalid_param1 + invalid_param2
-
-    with pytest.raises(ValueError):
-        register_custom_architecture()(
-            type(
-                "InvalidArch", (), {"_fit": _invalid_func, "_posterior": _invalid_func}
-            )
-        )
-
-    # Scenario: Class with valid `_fit` but invalid `_posterior` methods
-    def _valid_fit(self, searchspace, train_x, train_y):
-        return self and searchspace and train_x and train_y
-
-    with pytest.raises(ValueError):
-        register_custom_architecture()(
-            type("InvalidArch", (), {"_fit": _valid_fit, "_posterior": _invalid_func})
-        )
-
-    # Scenario: Both Valid
-    def _valid_posterior(self, candidates):
-        return self and candidates
-
-    register_custom_architecture()(
-        type("ValidArch", (), {"_fit": _valid_fit, "_posterior": _valid_posterior})
-    )
diff --git a/tests/test_deprecations.py b/tests/test_deprecations.py
index cfd2fff76..18be0a185 100644
--- a/tests/test_deprecations.py
+++ b/tests/test_deprecations.py
@@ -239,3 +239,11 @@ def test_deprecated_transform_interface(searchspace):
         searchspace.discrete.transform(
             data=searchspace.discrete.exp_rep, allow_extra=True
         )
+
+
+def test_deprecated_surrogate_registration():
+    """Using the deprecated registration mechanism raises a warning."""
+    from baybe.surrogates import register_custom_architecture
+
+    with pytest.raises(DeprecationError):
+        register_custom_architecture()

From 05ed5967645d1840e7cc52845abf0ac10e569b82 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 24 Jul 2024 16:41:40 +0200
Subject: [PATCH 42/97] Introduce SurrogateProtocol class to enable custom
 architectures

---
 CHANGELOG.md                             |  4 +++-
 baybe/acquisition/base.py                |  4 ++--
 baybe/recommenders/pure/bayesian/base.py |  4 ++--
 baybe/surrogates/base.py                 | 27 ++++++++++++++++++++----
 baybe/surrogates/custom.py               |  3 ++-
 5 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d379feb50..f7fb35d6c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -39,6 +39,7 @@ _ `_optional` subpackage for managing optional dependencies
   and `SubspaceContinuous` classes
 - Reworked mechanisms for surrogate input/output scaling configurable per class
 - `ParameterScalerProtocol` class for enabling user-defined input scaling mechanisms
+- `SurrogateProtocol` as an interface for user-defined surrogate architectures 
 
 ### Changed
 - Passing an `Objective` to `Campaign` is now optional
@@ -72,7 +73,8 @@ _ `_optional` subpackage for managing optional dependencies
 - Passing a dataframe via the `data` argument to the `transform` methods of
   `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous` is no longer possible.
   The dataframe must now be passed as positional argument.
-- A deprecation error is thrown when attempting to use `register_custom_architecture`
+- The role of `register_custom_architecture` has been taken over by
+  `baybe.surrogates.base.SurrogateProtocol`
 
 ## [0.9.1] - 2024-06-04
 ### Changed
diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index ef67a4572..eba6ab916 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -17,7 +17,7 @@
     unstructure_base,
 )
 from baybe.serialization.mixin import SerialMixin
-from baybe.surrogates.base import Surrogate
+from baybe.surrogates.base import SurrogateProtocol
 from baybe.utils.basic import classproperty, filter_attributes
 from baybe.utils.boolean import is_abstract
 from baybe.utils.dataframe import to_tensor
@@ -37,7 +37,7 @@ def is_mc(cls) -> bool:
 
     def to_botorch(
         self,
-        surrogate: Surrogate,
+        surrogate: SurrogateProtocol,
         searchspace: SearchSpace,
         measurements: pd.DataFrame,
     ):
diff --git a/baybe/recommenders/pure/bayesian/base.py b/baybe/recommenders/pure/bayesian/base.py
index cdc4f701a..e62db46f3 100644
--- a/baybe/recommenders/pure/bayesian/base.py
+++ b/baybe/recommenders/pure/bayesian/base.py
@@ -13,14 +13,14 @@
 from baybe.recommenders.pure.base import PureRecommender
 from baybe.searchspace import SearchSpace
 from baybe.surrogates import CustomONNXSurrogate, GaussianProcessSurrogate
-from baybe.surrogates.base import Surrogate
+from baybe.surrogates.base import SurrogateProtocol
 
 
 @define
 class BayesianRecommender(PureRecommender, ABC):
     """An abstract class for Bayesian Recommenders."""
 
-    surrogate_model: Surrogate = field(factory=GaussianProcessSurrogate)
+    surrogate_model: SurrogateProtocol = field(factory=GaussianProcessSurrogate)
     """The used surrogate model."""
 
     acquisition_function: AcquisitionFunction = field(
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 27622d121..524570718 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, Protocol
 
 import pandas as pd
 from attrs import define, field
@@ -62,8 +62,27 @@ class _NoTransform(Enum):
 """Sentinel to indicate the absence of a transform where `None` is ambiguous."""
 
 
+class SurrogateProtocol(Protocol):
+    """Type protocol specifying the interface surrogate models need to implement."""
+
+    def fit(
+        self,
+        searchspace: SearchSpace,
+        objective: Objective,
+        measurements: pd.DataFrame,
+    ) -> None:
+        """Fit the surrogate to training data in a given modelling context.
+
+        For details on the expected method arguments, see
+        :meth:`baybe.recommenders.base.RecommenderProtocol`.
+        """
+
+    def to_botorch(self) -> Model:
+        """Create the botorch-ready representation of the fitted model."""
+
+
 @define
-class Surrogate(ABC, SerialMixin):
+class Surrogate(ABC, SurrogateProtocol, SerialMixin):
     """Abstract base class for all surrogate models."""
 
     # Class variables
@@ -101,8 +120,8 @@ class Surrogate(ABC, SerialMixin):
     to make them ready for processing by the surrogate. Only available after the
     surrogate has been fitted."""
 
-    def to_botorch(self) -> Model:
-        """Create the botorch-ready representation of the model."""
+    def to_botorch(self) -> Model:  # noqa: D102
+        # See base class.
         from baybe.surrogates._adapter import AdapterModel
 
         return AdapterModel(self)
diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py
index 2b2086c75..248bf73a6 100644
--- a/baybe/surrogates/custom.py
+++ b/baybe/surrogates/custom.py
@@ -37,7 +37,8 @@ def register_custom_architecture(*args, **kwargs) -> NoReturn:
     """Deprecated! Raises an error when used."""  # noqa: D401
     raise DeprecationError(
         "The 'register_custom_architecture' decorator is no longer available. "
-        "There will be a replacement based on Python's `typing.Protocol` soon."
+        "Use :class:`baybe.surrogates.base.SurrogateProtocol` instead to define "
+        "your custom architectures."
     )
 
 

From 525aed3fb05914d8d84d274f6e756816f32f124a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 24 Jul 2024 16:54:16 +0200
Subject: [PATCH 43/97] Fix typo in method reference

---
 baybe/acquisition/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index eba6ab916..92cf44ee1 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -44,7 +44,7 @@ def to_botorch(
         """Create the botorch-ready representation of the function.
 
         The required structure of `measurements` is specified in
-        :meth:`babye.recommenders.base.RecommenderProtocol.recommend`.
+        :meth:`baybe.recommenders.base.RecommenderProtocol.recommend`.
         """
         import botorch.acquisition as botorch_acqf_module
 

From be5116351466c3292c00f425ec597832216eb911 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 24 Jul 2024 16:55:24 +0200
Subject: [PATCH 44/97] Refactor transformation steps in acquisition function
 translation

This solves two issues:
* It removes the dependency on the `SurrogateProtocol.transform_*`
  methods, getting us one step closer to a minimal protocol interface.
* It solves a design issue introduced in the last dev PR, where the
  transformations were extended such that they also include scaling.
  The transformations needed in the translation step should however
  **not** include scaling, which a purely surrogate internal step.
  In particular, we want the layout such that `best_f` and `X_baseline`
  can be provided in a representation that does not need to know
  anything about surrogate-internal mechanisms.
---
 baybe/acquisition/base.py                | 8 +++++---
 baybe/recommenders/pure/bayesian/base.py | 2 +-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index 92cf44ee1..464ec4da0 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -10,7 +10,8 @@
 import pandas as pd
 from attrs import define
 
-from baybe.searchspace import SearchSpace
+from baybe.objectives.base import Objective
+from baybe.searchspace.core import SearchSpace
 from baybe.serialization.core import (
     converter,
     get_base_structure_hook,
@@ -39,6 +40,7 @@ def to_botorch(
         self,
         surrogate: SurrogateProtocol,
         searchspace: SearchSpace,
+        objective: Objective,
         measurements: pd.DataFrame,
     ):
         """Create the botorch-ready representation of the function.
@@ -51,8 +53,8 @@ def to_botorch(
         acqf_cls = getattr(botorch_acqf_module, self.__class__.__name__)
         params_dict = filter_attributes(object=self, callable_=acqf_cls.__init__)
 
-        train_x = surrogate.transform_inputs(measurements)
-        train_y = surrogate.transform_outputs(measurements)
+        train_x = searchspace.transform(measurements)
+        train_y = objective.transform(measurements)
 
         signature_params = signature(acqf_cls).parameters
         additional_params = {}
diff --git a/baybe/recommenders/pure/bayesian/base.py b/baybe/recommenders/pure/bayesian/base.py
index e62db46f3..a0c2dfca2 100644
--- a/baybe/recommenders/pure/bayesian/base.py
+++ b/baybe/recommenders/pure/bayesian/base.py
@@ -52,7 +52,7 @@ def _setup_botorch_acqf(
         """Create the acquisition function for the current training data."""  # noqa: E501
         self.surrogate_model.fit(searchspace, objective, measurements)
         self._botorch_acqf = self.acquisition_function.to_botorch(
-            self.surrogate_model, searchspace, measurements
+            self.surrogate_model, searchspace, objective, measurements
         )
 
     def recommend(  # noqa: D102

From a1eacc0321a1c1ca8163fd8df7776050ec917c43 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 24 Jul 2024 17:08:11 +0200
Subject: [PATCH 45/97] Fix remaining surrogate-external transformation calls

---
 baybe/recommenders/pure/bayesian/botorch.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/baybe/recommenders/pure/bayesian/botorch.py b/baybe/recommenders/pure/bayesian/botorch.py
index f28b417bb..8f2d12642 100644
--- a/baybe/recommenders/pure/bayesian/botorch.py
+++ b/baybe/recommenders/pure/bayesian/botorch.py
@@ -104,7 +104,7 @@ def _recommend_discrete(
         from botorch.optim import optimize_acqf_discrete
 
         # determine the next set of points to be tested
-        candidates_comp = self.surrogate_model.transform_inputs(candidates_exp)
+        candidates_comp = subspace_discrete.transform(candidates_exp)
         points, _ = optimize_acqf_discrete(
             self._botorch_acqf, batch_size, to_tensor(candidates_comp)
         )
@@ -216,7 +216,7 @@ def _recommend_hybrid(
         from botorch.optim import optimize_acqf_mixed
 
         # Transform discrete candidates
-        candidates_comp = self.surrogate_model.transform_inputs(candidates_exp)
+        candidates_comp = searchspace.transform(candidates_exp)
 
         if len(candidates_comp) > 0:
             # Calculate the number of samples from the given percentage

From 53e4adefc90b9dd6ad3707c8a4d02533e4a76413 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 25 Jul 2024 11:44:10 +0200
Subject: [PATCH 46/97] Implement torch-based column transformer

---
 baybe/utils/scaling.py | 44 +++++++++++++++++++++++++++++++++
 tests/test_scaling.py  | 56 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)
 create mode 100644 tests/test_scaling.py

diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index 3c1b419cf..049960f13 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -2,9 +2,14 @@
 
 from __future__ import annotations
 
+import itertools
 from typing import Protocol
 
 import pandas as pd
+from attrs import define, field
+from attrs.validators import deep_iterable, deep_mapping, instance_of
+from botorch.models.transforms.input import InputTransform
+from torch import Tensor
 
 
 class ParameterScalerProtocol(Protocol):
@@ -20,3 +25,42 @@ def fit(self, df: pd.DataFrame, /) -> None:
 
     def transform(self, df: pd.DataFrame, /) -> pd.DataFrame:
         """Transform a parameter dataframe using the fitted scaling logic."""
+
+
+@define
+class ColumnTransformer:
+    """Class for applying transforms to individual columns of tensors."""
+
+    mapping: dict[tuple[int, ...], InputTransform] = field(
+        validator=deep_mapping(
+            mapping_validator=instance_of(dict),
+            key_validator=deep_iterable(
+                member_validator=instance_of(int), iterable_validator=instance_of(tuple)
+            ),
+            value_validator=instance_of(InputTransform),
+        )
+    )
+    """A mapping defining what transform to apply to which columns."""
+
+    @mapping.validator
+    def _validate_mapping(self, _, value: dict):
+        """Validate that the each column is assigned to at most one transformer."""
+        for x, y in itertools.combinations(value.keys(), 2):
+            if not set(x).isdisjoint(y):
+                raise ValueError(
+                    f"The provided column specifications {x} and {y} are not disjoint."
+                )
+
+    def fit(self, x: Tensor, /) -> None:
+        """Fit the transformer to the given tensor."""
+        for cols, transformer in self.mapping.items():
+            transformer.train()
+            transformer(x[..., cols])
+
+    def transform(self, x: Tensor, /) -> Tensor:
+        """Transform the given tensor."""
+        out = x.clone()
+        for cols, transformer in self.mapping.items():
+            transformer.eval()
+            out[..., cols] = transformer(out[..., cols])
+        return out
diff --git a/tests/test_scaling.py b/tests/test_scaling.py
new file mode 100644
index 000000000..506bf79b1
--- /dev/null
+++ b/tests/test_scaling.py
@@ -0,0 +1,56 @@
+"""Scaling tests."""
+
+import math
+from unittest.mock import Mock
+
+import pytest
+import torch
+from botorch.models.transforms.input import InputStandardize, InputTransform, Normalize
+
+from baybe.utils.scaling import ColumnTransformer
+
+
+def test_column_transformer():
+    """Basic test that validates the transformation of a ColumnTransformer."""
+    # Define test input
+    base = torch.stack([-1 * torch.ones(5), torch.ones(5)])
+    x_train = 20 * base
+    x_test = 10 * base
+
+    # Create and fit transformer
+    mapping = {
+        (0, 2): Normalize(2),
+        (1, 3): InputStandardize(2),
+    }
+    transformer = ColumnTransformer(mapping)
+    transformer.fit(x_train)
+
+    # Transform data
+    y_train = transformer.transform(x_train)
+    y_test = transformer.transform(x_test)
+
+    # Expected results
+    s = 1 / math.sqrt(2)
+    target_train = torch.tensor(
+        [
+            [0.0, -s, 0.0, -s, -20.0],
+            [1.0, s, 1.0, s, 20.0],
+        ]
+    )
+    target_test = torch.tensor(
+        [
+            [0.25, -s / 2, 0.25, -s / 2, -10.0],
+            [0.75, s / 2, 0.75, s / 2, 10.0],
+        ]
+    )
+
+    assert torch.allclose(y_train, target_train)
+    assert torch.allclose(y_test, target_test)
+
+
+def test_non_disjoint_column_mapping():
+    """Creating a column transformer with non-disjoint columns raises an error."""
+    t = Mock(InputTransform)
+    mapping = {(0, 1): t, (1, 2): t}
+    with pytest.raises(ValueError, match="are not disjoint"):
+        ColumnTransformer(mapping)

From 255d91ba44e5ec9a7e6c002aeb7c2177a0f53bb8 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 25 Jul 2024 15:40:23 +0200
Subject: [PATCH 47/97] Refactor scaling logic

The new layout now explictly exposes the three layers needed:
* User level: Dataframe in experimental representation to `Posterior`
* Botorch level: Tensor in computational representation to `Posterior`
* Surrogate level: Tensor in scaled computational representation to
  scaled `Posterior`

Particular benefits:
* Scaling is now completely encapsulated inside the surrogate so that
  objects outside do not need to bother about surrogate internals.
  That means, questions like "do we need to scale certain quantities
  before passing them to the surrogate" (such as `best_f` or
  `X_pending`) are trivially answered with "No", since scaling is not
  visible outside the surrogate.
* Scaling is now part of the computational torch graph, meaning
  that backpropagation through the entire surrogate model is
  supported.
---
 CHANGELOG.md                              |   1 -
 baybe/searchspace/core.py                 |   9 +
 baybe/surrogates/_adapter.py              |   2 +-
 baybe/surrogates/base.py                  | 197 +++++++++-------------
 baybe/surrogates/gaussian_process/core.py |  11 +-
 baybe/surrogates/ngboost.py               |   6 +-
 baybe/surrogates/random_forest.py         |   6 +-
 baybe/utils/scaling.py                    |  17 --
 8 files changed, 98 insertions(+), 151 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f7fb35d6c..829632cc6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -38,7 +38,6 @@ _ `_optional` subpackage for managing optional dependencies
 - `comp_rep_columns` property for `Parameter`, `SearchSpace`, `SubspaceDiscrete`
   and `SubspaceContinuous` classes
 - Reworked mechanisms for surrogate input/output scaling configurable per class
-- `ParameterScalerProtocol` class for enabling user-defined input scaling mechanisms
 - `SurrogateProtocol` as an interface for user-defined surrogate architectures 
 
 ### Changed
diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py
index 262f35f88..3b02bef00 100644
--- a/baybe/searchspace/core.py
+++ b/baybe/searchspace/core.py
@@ -290,6 +290,15 @@ def n_tasks(self) -> int:
         except StopIteration:
             return 1
 
+    def get_comp_rep_parameter_indices(self, name: str, /) -> tuple[int, ...]:
+        """Find a parameter's column indices in the computational representation."""
+        # TODO: The "startswith" approach is not ideal since it relies on the implicit
+        #   assumption that the substrings match. A more robust approach would
+        #   be to generate this mapping while building the comp rep.
+        return tuple(
+            i for i, col in enumerate(self.comp_rep_columns) if col.startswith(name)
+        )
+
     @staticmethod
     def estimate_product_space_size(parameters: Iterable[Parameter]) -> MemorySize:
         """Estimate an upper bound for the memory size of a product space.
diff --git a/baybe/surrogates/_adapter.py b/baybe/surrogates/_adapter.py
index a766446bd..6d167fe84 100644
--- a/baybe/surrogates/_adapter.py
+++ b/baybe/surrogates/_adapter.py
@@ -47,4 +47,4 @@ def posterior(  # noqa: D102
             raise NotImplementedError(
                 "The optional model posterior arguments are not yet implemented."
             )
-        return self._surrogate._posterior(X)
+        return self._surrogate._posterior_comp_rep(X)
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 524570718..55023ae77 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, ClassVar, Protocol
 
@@ -17,9 +16,7 @@
     UnstructuredValue,
     UnstructureHook,
 )
-from sklearn.preprocessing import MinMaxScaler
 
-from baybe.exceptions import ModelNotTrainedError
 from baybe.objectives.base import Objective
 from baybe.parameters.base import Parameter
 from baybe.searchspace import SearchSpace
@@ -30,13 +27,13 @@
 )
 from baybe.serialization.mixin import SerialMixin
 from baybe.utils.dataframe import to_tensor
-from baybe.utils.scaling import ParameterScalerProtocol
+from baybe.utils.scaling import ColumnTransformer
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
+    from botorch.models.transforms.input import InputTransform
     from botorch.models.transforms.outcome import OutcomeTransform
     from botorch.posteriors import GPyTorchPosterior, Posterior
-    from sklearn.compose import ColumnTransformer
     from torch import Tensor
 
 _ONNX_ENCODING = "latin-1"
@@ -93,32 +90,18 @@ class Surrogate(ABC, SurrogateProtocol, SerialMixin):
     """Class variable encoding whether or not the surrogate supports transfer
     learning."""
 
-    _input_transform: Callable[[pd.DataFrame], pd.DataFrame] | None = field(
-        init=False, default=None, eq=False
-    )
-    """Callable preparing surrogate inputs for training/prediction.
+    _searchspace: SearchSpace | None = field(init=False, default=None, eq=False)
+    """The search space on which the surrogate operates. Available after fitting."""
 
-    Transforms a dataframe containing parameter configurations in experimental
-    representation to a corresponding dataframe containing their **scaled**
-    computational representation. Only available after the surrogate has been fitted."""
-
-    _output_transform: Callable[[pd.DataFrame], pd.DataFrame] | None = field(
-        init=False, default=None, eq=False
-    )
-    """Callable preparing surrogate outputs for training.
-
-    Transforms a dataframe containing target measurements in experimental
-    representation to a corresponding dataframe containing their **scaled**
-    computational representation. Only available after the surrogate has been fitted."""
-
-    # TODO: type should be `Standardize | _NoTransform`` but is currently
-    #   omitted due to: https://github.com/python-attrs/cattrs/issues/531
+    # TODO: type should be
+    #   `botorch.models.transforms.outcome.Standardize | _NoTransform`
+    #   but is currently omitted due to:
+    #   https://github.com/python-attrs/cattrs/issues/531
     _output_scaler = field(init=False, default=None, eq=False)
-    """Optional callable for scaling output values.
+    """Scaler for transforming output values. Available after fitting.
 
     Scales a tensor containing target measurements in computational representation
-    to make them ready for processing by the surrogate. Only available after the
-    surrogate has been fitted."""
+    to make them digestible for the model-specific, scale-agnostic posterior logic."""
 
     def to_botorch(self) -> Model:  # noqa: D102
         # See base class.
@@ -127,11 +110,13 @@ def to_botorch(self) -> Model:  # noqa: D102
         return AdapterModel(self)
 
     @staticmethod
-    def _make_parameter_scaler(
+    def _make_parameter_scaler_factory(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | None:
-        """Return the scaler to be used for the given parameter."""
-        return MinMaxScaler()
+    ) -> type[InputTransform] | None:
+        """Return the scaler factory to be used for the given parameter."""
+        from botorch.models.transforms.input import Normalize
+
+        return Normalize
 
     @staticmethod
     def _make_target_scaler() -> OutcomeTransform | None:
@@ -143,19 +128,19 @@ def _make_target_scaler() -> OutcomeTransform | None:
 
     def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
         """Make the input scaler for transforming computational dataframes."""
-        from sklearn.compose import make_column_transformer
-
-        # Create the composite scaler from the parameter-wise scaler objects
-        transformers = [
-            (
-                "passthrough" if (s := self._make_parameter_scaler(p)) is None else s,
-                [c for c in p.comp_rep_columns if c in searchspace.comp_rep_columns],
-            )
-            for p in searchspace.parameters
-        ]
-        scaler = make_column_transformer(*transformers, verbose_feature_names_out=False)
-
-        scaler.fit(searchspace.comp_rep_bounds)
+        # Create a composite scaler from parameter-wise scaler objects
+        mapping: dict[tuple[int, ...], InputTransform] = {}
+        for p in searchspace.parameters:
+            idxs = searchspace.get_comp_rep_parameter_indices(p.name)
+            factory = self._make_parameter_scaler_factory(p)
+            if factory is None:
+                continue
+            transformer = factory(len(idxs))
+            mapping[idxs] = transformer
+        scaler = ColumnTransformer(mapping)
+
+        # Fit the scaler to the parameter bounds
+        scaler.fit(to_tensor(searchspace.comp_rep_bounds))
 
         return scaler
 
@@ -163,40 +148,60 @@ def _make_output_scaler(
         self, objective: Objective, measurements: pd.DataFrame
     ) -> OutcomeTransform | _NoTransform:
         """Make the output scaler for transforming computational dataframes."""
-        import torch
-
         scaler = self._make_target_scaler()
         if scaler is None:
             return _IDENTITY_TRANSFORM
 
         # TODO: Consider taking into account target boundaries when available
-        scaler(torch.from_numpy(objective.transform(measurements).values))
+        scaler(to_tensor(objective.transform(measurements)))
         scaler.eval()
 
         return scaler
 
-    def transform_inputs(self, df: pd.DataFrame, /) -> pd.DataFrame:
-        """Transform an experimental parameter dataframe."""
-        if self._input_transform is None:
-            raise ModelNotTrainedError("The model must be trained first.")
-        return self._input_transform(df)
+    def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
+        """Compute the posterior for candidates in experimental representation.
 
-    def transform_outputs(self, df: pd.DataFrame, /) -> pd.DataFrame:
-        """Transform an experimental measurement dataframe."""
-        if self._output_transform is None:
-            raise ModelNotTrainedError("The model must be trained first.")
-        return self._output_transform(df)
+        Takes a dataframe of parameter configurations in **experimental representation**
+        and returns the corresponding posterior object. Therefore, the method serves as
+        the user-facing entry point for accessing model predictions.
+        """
+        return self._posterior_comp_rep(
+            to_tensor(self._searchspace.transform(candidates))
+        )
 
-    def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
-        """Evaluate the surrogate model at the given candidate points."""
-        p = self._posterior(to_tensor(self.transform_inputs(candidates)))
+    def _posterior_comp_rep(self, candidates: Tensor, /) -> Posterior:
+        """Compute the posterior for candidates in computational representation.
+
+        Takes a tensor of parameter configurations in **computational representation**
+        and returns the corresponding posterior object. Therefore, the method provides
+        the entry point for queries coming from computational layers, for instance,
+        BoTorch's `optimize_*` functions.
+        """
+        p = self._posterior(self._input_scaler.transform(candidates))
         if self._output_scaler is not _IDENTITY_TRANSFORM:
             p = self._output_scaler.untransform_posterior(p)
         return p
 
     @abstractmethod
     def _posterior(self, candidates: Tensor, /) -> Posterior:
-        """Perform the actual posterior evaluation logic."""
+        """Perform the actual model-specific posterior evaluation logic.
+
+        This method is supposed to be overridden by subclasses to implement their
+        model-specific surrogate architecture. Internally, the method is called by the
+        base class with a **scaled** tensor of candidates in **computational
+        representation**, where the scaling is configurable by the subclass via its
+        other methods. The base class also takes care of transforming the returned
+        posterior back to the original scale according to the defined scalers.
+
+        This means:
+        -----------
+        Subclasses implementing this method do not have to bother about
+        pre-/postprocessing of the in-/output. Instead, they only need to implement the
+        mathematical operation of computing the posterior for the given input according
+        to their model specifications and can implicitly that scaling is handled
+        appropriately outside. In short: the returned posterior simply needs to be on
+        the same scale as the given input.
+        """
 
     @staticmethod
     def _get_model_context(searchspace: SearchSpace, objective: Objective) -> Any:
@@ -242,70 +247,20 @@ def fit(
                 "Continuous search spaces are currently only supported by GPs."
             )
 
-        # Create scaler objects
-        input_scaler = self._make_input_scaler(searchspace)
-        output_scaler = self._make_output_scaler(objective, measurements)
-
-        def transform_inputs(df: pd.DataFrame, /) -> pd.DataFrame:
-            """Fitted input transformation pipeline."""
-            # IMPROVE: This method currently relies on two workarounds required
-            #   due the working mechanism of sklearn's ColumnTransformer:
-            #   * Unfortunately, the transformer returns a raw array, meaning that
-            #       column names need to be manually attached afterward.
-            #   * In certain cases (e.g. in hybrid spaces), the method needs
-            #       to transform only a subset of columns. Unfortunately, it is not
-            #       possible to use a subset of columns once the transformer is set up,
-            #       which is a side-effect of the first point. As a workaround,
-            #       we thus fill the missing columns with NaN and subselect afterward.
-
-            # For the workaround, collect all comp rep columns of the parameters
-            # that are actually present in the given dataframe. At the end,
-            # we'll filter the transformed augmented dataframe down to these columns.
-            exp_rep_cols = [p.name for p in searchspace.parameters]
-            comp_rep_cols = []
-            for col in [c for c in df.columns if c in exp_rep_cols]:
-                parameter = next(p for p in searchspace.parameters if p.name == col)
-                comp_rep_cols.extend(parameter.comp_rep_columns)
-
-            # Actual workaround: augment the dataframe with NaN for missing parameters
-            df_augmented = df.reindex(columns=exp_rep_cols)
-
-            # The actual transformation step
-            out = input_scaler.transform(
-                searchspace.transform(df_augmented, allow_extra=True)
-            )
-            out = pd.DataFrame(
-                out, index=df.index, columns=input_scaler.get_feature_names_out()
-            )
-
-            # Undo the augmentation, taking into account that not all comp rep
-            # parameter columns may actually be part of the search space due
-            # to other preprocessing steps.
-            comp_rep_cols = list(set(comp_rep_cols).intersection(out.columns))
-            return out[comp_rep_cols]
-
-        def transform_outputs(df: pd.DataFrame, /) -> pd.DataFrame:
-            """Fitted output transformation pipeline."""
-            import torch
-
-            dft = objective.transform(df)
-
-            if output_scaler is _IDENTITY_TRANSFORM:
-                return dft
-
-            out = output_scaler(torch.from_numpy(dft.values))[0]
-            return pd.DataFrame(out.numpy(), index=df.index, columns=dft.columns)
+        # Remember on which search space the model is trained
+        self._searchspace = searchspace
 
-        # Store context-specific transformations
-        self._input_transform = transform_inputs
-        self._output_transform = transform_outputs
-        self._output_scaler = output_scaler
+        # Create context-specific transformations
+        self._input_scaler = self._make_input_scaler(searchspace)
+        self._output_scaler = self._make_output_scaler(objective, measurements)
 
         # Transform and fit
-        train_x, train_y = to_tensor(
-            self.transform_inputs(measurements),
-            self.transform_outputs(measurements),
+        train_x_comp_rep, train_y_comp_rep = to_tensor(
+            searchspace.transform(measurements),
+            objective.transform(measurements),
         )
+        train_x = self._input_scaler.transform(train_x_comp_rep)
+        train_y = self._output_scaler(train_y_comp_rep)[0]
         self._fit(train_x, train_y, self._get_model_context(searchspace, objective))
 
     @abstractmethod
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 5b00982e6..147a4fc59 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -6,7 +6,6 @@
 
 from attrs import define, field
 from attrs.validators import instance_of
-from sklearn.preprocessing import MinMaxScaler
 
 from baybe.objective import Objective
 from baybe.parameters.base import Parameter
@@ -25,10 +24,10 @@
     DefaultKernelFactory,
     _default_noise_factory,
 )
-from baybe.utils.scaling import ParameterScalerProtocol
 
 if TYPE_CHECKING:
     from botorch.models.model import Model
+    from botorch.models.transforms.input import InputTransform
     from botorch.posteriors import Posterior
     from torch import Tensor
 
@@ -113,16 +112,18 @@ def to_botorch(self) -> Model:  # noqa: D102
         return self._model
 
     @staticmethod
-    def _make_parameter_scaler(
+    def _make_parameter_scaler_factory(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | None:
+    ) -> type[InputTransform] | None:
         # See base class.
 
         # Task parameters are handled separately through an index kernel
         if isinstance(parameter, TaskParameter):
             return
 
-        return MinMaxScaler()
+        from botorch.models.transforms.input import Normalize
+
+        return Normalize
 
     @staticmethod
     def _get_model_context(
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index b0ce3849c..b157ce65c 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -18,9 +18,9 @@
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
-from baybe.utils.scaling import ParameterScalerProtocol
 
 if TYPE_CHECKING:
+    from botorch.models.transforms.input import InputTransform
     from torch import Tensor
 
 
@@ -54,9 +54,9 @@ def __attrs_post_init__(self):
         self.model_params = {**self._default_model_params, **self.model_params}
 
     @staticmethod
-    def _make_parameter_scaler(
+    def _make_parameter_scaler_factory(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | None:
+    ) -> type[InputTransform] | None:
         # See base class.
 
         # Tree-like models do not require any input scaling
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 776a9a830..07466c03f 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -19,9 +19,9 @@
 from baybe.surrogates.base import GaussianSurrogate
 from baybe.surrogates.utils import batchify, catch_constant_targets
 from baybe.surrogates.validation import get_model_params_validator
-from baybe.utils.scaling import ParameterScalerProtocol
 
 if TYPE_CHECKING:
+    from botorch.models.transforms.input import InputTransform
     from torch import Tensor
 
 
@@ -49,9 +49,9 @@ class RandomForestSurrogate(GaussianSurrogate):
     """The actual model."""
 
     @staticmethod
-    def _make_parameter_scaler(
+    def _make_parameter_scaler_factory(
         parameter: Parameter,
-    ) -> ParameterScalerProtocol | None:
+    ) -> type[InputTransform] | None:
         # See base class.
 
         # Tree-like models do not require any input scaling
diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index 049960f13..f7fd6fca7 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -3,30 +3,13 @@
 from __future__ import annotations
 
 import itertools
-from typing import Protocol
 
-import pandas as pd
 from attrs import define, field
 from attrs.validators import deep_iterable, deep_mapping, instance_of
 from botorch.models.transforms.input import InputTransform
 from torch import Tensor
 
 
-class ParameterScalerProtocol(Protocol):
-    """Type protocol specifying the interface parameter scalers need to implement.
-
-    The protocol is compatible with sklearn scalers such as
-    :class:`sklearn.preprocessing.MinMaxScaler` or
-    :class:`sklearn.preprocessing.MaxAbsScaler`.
-    """
-
-    def fit(self, df: pd.DataFrame, /) -> None:
-        """Fit the scaler to a given dataframe containing parameter configurations."""
-
-    def transform(self, df: pd.DataFrame, /) -> pd.DataFrame:
-        """Transform a parameter dataframe using the fitted scaling logic."""
-
-
 @define
 class ColumnTransformer:
     """Class for applying transforms to individual columns of tensors."""

From 47e68195e930c2ea5b94f720d308b8bae9ec4fae Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 26 Jul 2024 09:41:21 +0200
Subject: [PATCH 48/97] Add missing transform flags

---
 baybe/acquisition/base.py                   | 2 +-
 baybe/recommenders/pure/bayesian/botorch.py | 2 +-
 baybe/surrogates/base.py                    | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/baybe/acquisition/base.py b/baybe/acquisition/base.py
index 464ec4da0..eadb3dc75 100644
--- a/baybe/acquisition/base.py
+++ b/baybe/acquisition/base.py
@@ -53,7 +53,7 @@ def to_botorch(
         acqf_cls = getattr(botorch_acqf_module, self.__class__.__name__)
         params_dict = filter_attributes(object=self, callable_=acqf_cls.__init__)
 
-        train_x = searchspace.transform(measurements)
+        train_x = searchspace.transform(measurements, allow_extra=True)
         train_y = objective.transform(measurements)
 
         signature_params = signature(acqf_cls).parameters
diff --git a/baybe/recommenders/pure/bayesian/botorch.py b/baybe/recommenders/pure/bayesian/botorch.py
index 8f2d12642..de05cad7d 100644
--- a/baybe/recommenders/pure/bayesian/botorch.py
+++ b/baybe/recommenders/pure/bayesian/botorch.py
@@ -216,7 +216,7 @@ def _recommend_hybrid(
         from botorch.optim import optimize_acqf_mixed
 
         # Transform discrete candidates
-        candidates_comp = searchspace.transform(candidates_exp)
+        candidates_comp = searchspace.transform(candidates_exp, allow_missing=True)
 
         if len(candidates_comp) > 0:
             # Calculate the number of samples from the given percentage
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 55023ae77..eb93caa68 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -256,7 +256,7 @@ def fit(
 
         # Transform and fit
         train_x_comp_rep, train_y_comp_rep = to_tensor(
-            searchspace.transform(measurements),
+            searchspace.transform(measurements, allow_extra=True),
             objective.transform(measurements),
         )
         train_x = self._input_scaler.transform(train_x_comp_rep)

From ee6439c73d7852ece35619f98e11a0eb619437ff Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 26 Jul 2024 09:41:39 +0200
Subject: [PATCH 49/97] Register de-/serialization hooks for SurrogateProtocol

---
 baybe/surrogates/base.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index eb93caa68..769ff450c 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -353,14 +353,17 @@ def wrapper(obj: StructuredValue) -> UnstructuredValue:
 #   class, which would avoid the nested wrapping below. However, this requires
 #   adjusting the base class (un-)structure hooks such that they consistently apply
 #   existing hooks of the concrete subclasses.
-converter.register_unstructure_hook(
-    Surrogate,
-    _make_hook_decode_onnx_str(
-        _block_serialize_custom_architecture(
-            lambda x: unstructure_base(x, overrides={"_model": override(omit=True)})
-        )
-    ),
+_unstructure_hook = _make_hook_decode_onnx_str(
+    _block_serialize_custom_architecture(
+        lambda x: unstructure_base(x, overrides={"_model": override(omit=True)})
+    )
 )
+converter.register_unstructure_hook(Surrogate, _unstructure_hook)
 converter.register_structure_hook(
     Surrogate, _make_hook_encode_onnx_str(get_base_structure_hook(Surrogate))
 )
+converter.register_unstructure_hook(SurrogateProtocol, _unstructure_hook)
+converter.register_structure_hook(
+    SurrogateProtocol,
+    _make_hook_encode_onnx_str(get_base_structure_hook(SurrogateProtocol)),
+)

From befd68f7739baefd108e02e3c786f297171a58de Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 26 Jul 2024 10:21:35 +0200
Subject: [PATCH 50/97] Add details on the requirements imposed by the
 surrogate protocol

---
 baybe/surrogates/base.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 769ff450c..244a82f54 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -75,7 +75,13 @@ def fit(
         """
 
     def to_botorch(self) -> Model:
-        """Create the botorch-ready representation of the fitted model."""
+        """Create the botorch-ready representation of the fitted model.
+
+        The :class:`botorch.models.model.Model` created by this method needs to be
+        configured such that it can be called with candidate points in **computational
+        representation**, that is, input of the form as obtained via
+        :meth:`baybe.searchspace.core.SearchSpace.transform`.
+        """
 
 
 @define

From 90fd62d71daa31442c53c4ac1c5fab826842ffcb Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 26 Jul 2024 11:50:16 +0200
Subject: [PATCH 51/97] Lazify validation of column transformer

---
 baybe/utils/scaling.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index f7fd6fca7..a841e8ff3 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -3,30 +3,39 @@
 from __future__ import annotations
 
 import itertools
+from typing import TYPE_CHECKING
 
 from attrs import define, field
 from attrs.validators import deep_iterable, deep_mapping, instance_of
-from botorch.models.transforms.input import InputTransform
-from torch import Tensor
+
+if TYPE_CHECKING:
+    from botorch.models.transforms.input import InputTransform
+    from torch import Tensor
 
 
 @define
 class ColumnTransformer:
     """Class for applying transforms to individual columns of tensors."""
 
-    mapping: dict[tuple[int, ...], InputTransform] = field(
-        validator=deep_mapping(
+    mapping: dict[tuple[int, ...], InputTransform] = field()
+    """A mapping defining what transform to apply to which columns."""
+
+    @mapping.validator
+    def _validate_mapping_types_lazily(self, attr, value):
+        """Perform transform isinstance check using lazy import."""
+        from botorch.models.transforms.input import InputTransform
+
+        validator = deep_mapping(
             mapping_validator=instance_of(dict),
             key_validator=deep_iterable(
                 member_validator=instance_of(int), iterable_validator=instance_of(tuple)
             ),
             value_validator=instance_of(InputTransform),
         )
-    )
-    """A mapping defining what transform to apply to which columns."""
+        validator(self, attr, value)
 
     @mapping.validator
-    def _validate_mapping(self, _, value: dict):
+    def _validate_mapping_is_disjoint(self, _, value: dict):
         """Validate that the each column is assigned to at most one transformer."""
         for x, y in itertools.combinations(value.keys(), 2):
             if not set(x).isdisjoint(y):

From a84dd8ea0484470b7a35818f88e056c609ae6979 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 11:29:36 +0200
Subject: [PATCH 52/97] Optimize input scaling logic using walrus

---
 baybe/surrogates/base.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 244a82f54..4bd05521a 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -137,10 +137,9 @@ def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
         # Create a composite scaler from parameter-wise scaler objects
         mapping: dict[tuple[int, ...], InputTransform] = {}
         for p in searchspace.parameters:
-            idxs = searchspace.get_comp_rep_parameter_indices(p.name)
-            factory = self._make_parameter_scaler_factory(p)
-            if factory is None:
+            if (factory := self._make_parameter_scaler_factory(p)) is None:
                 continue
+            idxs = searchspace.get_comp_rep_parameter_indices(p.name)
             transformer = factory(len(idxs))
             mapping[idxs] = transformer
         scaler = ColumnTransformer(mapping)

From c12a99299213fbed51d3ce86096ed954864c28f3 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 11:34:45 +0200
Subject: [PATCH 53/97] Rephrase ColumnTransformer docstring

---
 baybe/utils/scaling.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index a841e8ff3..78fdf73f7 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -15,7 +15,7 @@
 
 @define
 class ColumnTransformer:
-    """Class for applying transforms to individual columns of tensors."""
+    """Class for applying separate transforms to different column groups of tensors."""
 
     mapping: dict[tuple[int, ...], InputTransform] = field()
     """A mapping defining what transform to apply to which columns."""

From 629f742ca2937bc7dd1091f24670d958f0b5d3c9 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 12:04:54 +0200
Subject: [PATCH 54/97] Avoid allow_missing=True by accessing discrete subspace

---
 baybe/recommenders/pure/bayesian/botorch.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baybe/recommenders/pure/bayesian/botorch.py b/baybe/recommenders/pure/bayesian/botorch.py
index de05cad7d..eac1b3b19 100644
--- a/baybe/recommenders/pure/bayesian/botorch.py
+++ b/baybe/recommenders/pure/bayesian/botorch.py
@@ -216,7 +216,7 @@ def _recommend_hybrid(
         from botorch.optim import optimize_acqf_mixed
 
         # Transform discrete candidates
-        candidates_comp = searchspace.transform(candidates_exp, allow_missing=True)
+        candidates_comp = searchspace.discrete.transform(candidates_exp)
 
         if len(candidates_comp) > 0:
             # Calculate the number of samples from the given percentage

From 179fa5aa88b644156129f679d6ab1f01391fb968 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 13:45:37 +0200
Subject: [PATCH 55/97] Make target scaler method return a factory

---
 baybe/surrogates/base.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 4bd05521a..fefb0cb6f 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -125,12 +125,11 @@ def _make_parameter_scaler_factory(
         return Normalize
 
     @staticmethod
-    def _make_target_scaler() -> OutcomeTransform | None:
-        """Return the scaler to be used for target scaling."""
+    def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
+        """Return the scaler factory to be used for target scaling."""
         from botorch.models.transforms.outcome import Standardize
 
-        # TODO: Multi-target extension
-        return Standardize(1)
+        return Standardize
 
     def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
         """Make the input scaler for transforming computational dataframes."""
@@ -153,10 +152,12 @@ def _make_output_scaler(
         self, objective: Objective, measurements: pd.DataFrame
     ) -> OutcomeTransform | _NoTransform:
         """Make the output scaler for transforming computational dataframes."""
-        scaler = self._make_target_scaler()
-        if scaler is None:
+        if (factory := self._make_target_scaler_factory()) is None:
             return _IDENTITY_TRANSFORM
 
+        # TODO: Multi-target extension
+        scaler = factory(1)
+
         # TODO: Consider taking into account target boundaries when available
         scaler(to_tensor(objective.transform(measurements)))
         scaler.eval()

From 8168c014b201497bb2032d0b87d2ad38aa3f6556 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 13:47:49 +0200
Subject: [PATCH 56/97] Mention customization in scaler method docstrings

---
 baybe/surrogates/base.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index fefb0cb6f..75f0f2c02 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -119,14 +119,22 @@ def to_botorch(self) -> Model:  # noqa: D102
     def _make_parameter_scaler_factory(
         parameter: Parameter,
     ) -> type[InputTransform] | None:
-        """Return the scaler factory to be used for the given parameter."""
+        """Return the scaler factory to be used for the given parameter.
+
+        This method is supposed to be overridden by subclasses to implement their
+        custom parameter scaling logic.
+        """
         from botorch.models.transforms.input import Normalize
 
         return Normalize
 
     @staticmethod
     def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
-        """Return the scaler factory to be used for target scaling."""
+        """Return the scaler factory to be used for target scaling.
+
+        This method is supposed to be overridden by subclasses to implement their
+        custom target scaling logic.
+        """
         from botorch.models.transforms.outcome import Standardize
 
         return Standardize

From b1ec8d35ad989cc6593d8a32e110cea1eb6b3634 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 14:07:20 +0200
Subject: [PATCH 57/97] Add validation to column index extraction method

---
 baybe/searchspace/core.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py
index 3b02bef00..e55542c90 100644
--- a/baybe/searchspace/core.py
+++ b/baybe/searchspace/core.py
@@ -292,6 +292,11 @@ def n_tasks(self) -> int:
 
     def get_comp_rep_parameter_indices(self, name: str, /) -> tuple[int, ...]:
         """Find a parameter's column indices in the computational representation."""
+        if name not in (p.name for p in self.parameters):
+            raise ValueError(
+                f"There exists no parameter named '{name}' in the search space."
+            )
+
         # TODO: The "startswith" approach is not ideal since it relies on the implicit
         #   assumption that the substrings match. A more robust approach would
         #   be to generate this mapping while building the comp rep.

From b83f39c9f88ce4ec93e2c38c89bd1f9c087ebeaf Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 14:22:48 +0200
Subject: [PATCH 58/97] Expand docstring of get_comp_rep_parameter_indices

---
 baybe/searchspace/core.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py
index e55542c90..31e1127ee 100644
--- a/baybe/searchspace/core.py
+++ b/baybe/searchspace/core.py
@@ -291,7 +291,19 @@ def n_tasks(self) -> int:
             return 1
 
     def get_comp_rep_parameter_indices(self, name: str, /) -> tuple[int, ...]:
-        """Find a parameter's column indices in the computational representation."""
+        """Find a parameter's column indices in the computational representation.
+
+        Args:
+            name: The name of the parameter whose columns indices are to be retrieved.
+
+        Raises:
+            ValueError: If no parameter with the provided name exists.
+
+        Returns:
+            A tuple containing the integer indices of the columns in the computational
+            representation associated with the parameter. When the parameter is not part
+            of the computational representation, an empty tuple is returned.
+        """
         if name not in (p.name for p in self.parameters):
             raise ValueError(
                 f"There exists no parameter named '{name}' in the search space."

From aafe2f7355102e5fd1f0a51f004d42759b9207fb Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 16:44:19 +0200
Subject: [PATCH 59/97] Add TODO note

---
 baybe/acquisition/acqfs.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/baybe/acquisition/acqfs.py b/baybe/acquisition/acqfs.py
index c4b74e94b..ad6e1d2e8 100644
--- a/baybe/acquisition/acqfs.py
+++ b/baybe/acquisition/acqfs.py
@@ -81,6 +81,8 @@ def get_integration_points(self, searchspace: SearchSpace) -> pd.DataFrame:
             ValueError: If the search space is purely continuous and
                 'sampling_n_points' was not provided.
         """
+        # TODO: Move the core logic to `SearchSpace` and ``Subspace*`` classes
+
         sampled_parts: list[pd.DataFrame] = []
         n_candidates: int | None = None
 

From 6dabe87722b28285800a3e4321c37eaa2e875460 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 16:55:17 +0200
Subject: [PATCH 60/97] Explicitly handle train/eval mode in ColumnTransformer

---
 baybe/utils/scaling.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index 78fdf73f7..4059e2402 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -20,6 +20,9 @@ class ColumnTransformer:
     mapping: dict[tuple[int, ...], InputTransform] = field()
     """A mapping defining what transform to apply to which columns."""
 
+    _is_trained: bool = field(default=False, init=False)
+    """Boolean indicating if the transformer has been trained."""
+
     @mapping.validator
     def _validate_mapping_types_lazily(self, attr, value):
         """Perform transform isinstance check using lazy import."""
@@ -45,14 +48,24 @@ def _validate_mapping_is_disjoint(self, _, value: dict):
 
     def fit(self, x: Tensor, /) -> None:
         """Fit the transformer to the given tensor."""
+        # Explicitly set flag to False to guarantee a clean state in case of
+        # exceptions occurring in the for-loop below
+        self._is_trained = False
+
         for cols, transformer in self.mapping.items():
             transformer.train()
             transformer(x[..., cols])
+            transformer.eval()
+        self._is_trained = True
 
     def transform(self, x: Tensor, /) -> Tensor:
         """Transform the given tensor."""
+        if not self._is_trained:
+            raise RuntimeError(
+                f"The {self.__class__.__name__} must be trained before it can be used."
+            )
+
         out = x.clone()
         for cols, transformer in self.mapping.items():
-            transformer.eval()
             out[..., cols] = transformer(out[..., cols])
         return out

From 57ee8b3312313ff49723bbf989d99a4a09a71c1d Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 21:43:41 +0200
Subject: [PATCH 61/97] Add docstring sections to posterior methods

---
 baybe/surrogates/base.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 75f0f2c02..5e1767dcb 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -178,6 +178,17 @@ def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
         Takes a dataframe of parameter configurations in **experimental representation**
         and returns the corresponding posterior object. Therefore, the method serves as
         the user-facing entry point for accessing model predictions.
+
+        Args:
+            candidates: A dataframe containing parameter configurations in
+                **experimental representation**.
+
+        Returns:
+            A :class:`botorch.posteriors.Posterior` object representing the posterior
+            distribution at the given candidate points, where the posterior is also
+            described in **experimental representation**. That is, the posterior values
+            lie in the same domain as the modelled targets/objective on which the
+            surrogate was trained via :meth:`baybe.surrogates.base.Surrogate.fit`.
         """
         return self._posterior_comp_rep(
             to_tensor(self._searchspace.transform(candidates))
@@ -190,6 +201,14 @@ def _posterior_comp_rep(self, candidates: Tensor, /) -> Posterior:
         and returns the corresponding posterior object. Therefore, the method provides
         the entry point for queries coming from computational layers, for instance,
         BoTorch's `optimize_*` functions.
+
+        Args:
+            candidates: A tensor containing parameter configurations in **computational
+                representation**.
+
+        Returns:
+            The same :class:`botorch.posteriors.Posterior` object as returned via
+            :meth:`baybe.surrogates.base.Surrogate.posterior`.
         """
         p = self._posterior(self._input_scaler.transform(candidates))
         if self._output_scaler is not _IDENTITY_TRANSFORM:
@@ -215,6 +234,17 @@ def _posterior(self, candidates: Tensor, /) -> Posterior:
         to their model specifications and can implicitly that scaling is handled
         appropriately outside. In short: the returned posterior simply needs to be on
         the same scale as the given input.
+
+        Args:
+            candidates: A tensor containing **pre-scaled** parameter configurations
+                in **computational representation**, as defined through the input scaler
+                obtained via :meth:`baybe.surrogates.base.Surrogate._make_input_scaler`.
+
+        Returns:
+            A :class:`botorch.posteriors.Posterior` object representing the
+            **scale-transformed** posterior distributions at the given candidate points,
+            where the posterior is described on the scale dictated by the output scaler
+            obtained via :meth:`baybe.surrogates.base.Surrogate._make_output_scaler`.
         """
 
     @staticmethod

From bcbcf3fa250d45a631a33aee7c7e3ef6838a374e Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 7 Aug 2024 14:56:50 +0200
Subject: [PATCH 62/97] Account for potentially non-existing output scaler

---
 baybe/surrogates/base.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 5e1767dcb..d26777cfb 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -304,7 +304,11 @@ def fit(
             objective.transform(measurements),
         )
         train_x = self._input_scaler.transform(train_x_comp_rep)
-        train_y = self._output_scaler(train_y_comp_rep)[0]
+        train_y = (
+            train_y_comp_rep
+            if self._output_scaler is _IDENTITY_TRANSFORM
+            else self._output_scaler(train_y_comp_rep)[0]
+        )
         self._fit(train_x, train_y, self._get_model_context(searchspace, objective))
 
     @abstractmethod

From 5e63cb7bc776360eecb4c056b0112cc48c5c0e30 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 31 Jul 2024 21:59:35 +0200
Subject: [PATCH 63/97] Fix bug in GP posterior computation

The `GaussianProcessSurrogate._posterior` method implements the
posterior computation where the input is already scale-transforemd.
This should **not** be passed to botorch, because the transition
point to botorch is on a layer where the concept of scaling
is irrelevant / does not exist (scaling is a purely surrogate-
internal process). Instead, we go via the adapter model, which wraps
the appropriate method.
---
 baybe/surrogates/gaussian_process/core.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 147a4fc59..52dfb36cd 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -26,7 +26,6 @@
 )
 
 if TYPE_CHECKING:
-    from botorch.models.model import Model
     from botorch.models.transforms.input import InputTransform
     from botorch.posteriors import Posterior
     from torch import Tensor
@@ -106,11 +105,6 @@ def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate:
         """Create a Gaussian process surrogate from one of the defined presets."""
         return make_gp_from_preset(preset)
 
-    def to_botorch(self) -> Model:  # noqa: D102
-        # See base class.
-
-        return self._model
-
     @staticmethod
     def _make_parameter_scaler_factory(
         parameter: Parameter,

From ab26cdd07057fe5e878146c700071b0ca89eaecd Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 7 Aug 2024 15:01:56 +0200
Subject: [PATCH 64/97] Change plan and expose internal GP, but with disabled
 base class scaling

The problem is that qNIPV requires a model that can `fantasize`, which
our `AdapterModel` is not capable of. We thus expose the actual botorch
GP model instead. But for this to work correctly, we still need to
ensure that scaling happens **after** the transition point to
`optimize_acqf_*`:

`optimize_acqf_*` is called with **unscaled** parameters in **computational
representation**, so it must be provided a surrogate that operates on
the same domain. In regular operation, the `AdapterModel` takes care
of this by exposing the method `posterior_comp_rep` as the correct entry
point. When we now expose the internal botorch GP model, the exposed
posterior method was calibrated on the already scaled training data,
resulting in a mismatch. The workaround is to deactive the base class
scaling and instead let the GP internally do the scaling work.
---
 baybe/surrogates/gaussian_process/core.py | 41 +++++++++++++++++++----
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 52dfb36cd..5e9b22350 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -9,7 +9,6 @@
 
 from baybe.objective import Objective
 from baybe.parameters.base import Parameter
-from baybe.parameters.categorical import TaskParameter
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.base import Surrogate
 from baybe.surrogates.gaussian_process.kernel_factory import (
@@ -26,7 +25,9 @@
 )
 
 if TYPE_CHECKING:
+    from botorch.models.model import Model
     from botorch.models.transforms.input import InputTransform
+    from botorch.models.transforms.outcome import OutcomeTransform
     from botorch.posteriors import Posterior
     from torch import Tensor
 
@@ -75,6 +76,19 @@ def get_numerical_indices(self, n_inputs: int) -> list[int]:
 class GaussianProcessSurrogate(Surrogate):
     """A Gaussian process surrogate model."""
 
+    # Note [Scaling Workaround]
+    # -------------------------
+    # For GPs, we deactivate the base class scaling and instead let the botorch
+    # model internally handle input/output scaling. The reasons is that we need to
+    # make `to_botorch` expose the actual botorch GP object, instead of going
+    # via the `AdapterModel`, because certain acquisition functions (like qNIPV)
+    # require the capability to `fantasize`, which the `AdapterModel` does not support.
+    # The base class scaling thus needs to be disabled since otherwise the botorch GP
+    # object would be trained on pre-scaled input/output data. This would cause a
+    # problem since the resulting `posterior` method of that object is exposed
+    # to `optimize_acqf_*`, which is configured to be called on the original scale.
+    # Moving the scaling operation into the botorch GP object avoids this conflict.
+
     # Class variables
     joint_posterior: ClassVar[bool] = True
     # See base class.
@@ -105,19 +119,26 @@ def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate:
         """Create a Gaussian process surrogate from one of the defined presets."""
         return make_gp_from_preset(preset)
 
+    def to_botorch(self) -> Model:  # noqa: D102
+        # See base class.
+
+        return self._model
+
     @staticmethod
     def _make_parameter_scaler_factory(
         parameter: Parameter,
     ) -> type[InputTransform] | None:
         # See base class.
 
-        # Task parameters are handled separately through an index kernel
-        if isinstance(parameter, TaskParameter):
-            return
+        # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
+        return
 
-        from botorch.models.transforms.input import Normalize
+    @staticmethod
+    def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
+        # See base class.
 
-        return Normalize
+        # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
+        return
 
     @staticmethod
     def _get_model_context(
@@ -139,6 +160,12 @@ def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None
 
         numerical_idxs = context.get_numerical_indices(train_x.shape[-1])
 
+        # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
+        input_transform = botorch.models.transforms.Normalize(
+            train_x.shape[-1], bounds=context.parameter_bounds, indices=numerical_idxs
+        )
+        outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1])
+
         # extract the batch shape of the training data
         batch_shape = train_x.shape[:-2]
 
@@ -176,6 +203,8 @@ def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None
         self._model = botorch.models.SingleTaskGP(
             train_x,
             train_y,
+            input_transform=input_transform,
+            outcome_transform=outcome_transform,
             mean_module=mean_module,
             covar_module=covar_module,
             likelihood=likelihood,

From 06a99b98f6c58203635ed67e65b24f292509568a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 7 Aug 2024 15:46:04 +0200
Subject: [PATCH 65/97] Activate mypy for surrogates

---
 mypy.ini | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mypy.ini b/mypy.ini
index 10b678f24..51a9b8bb3 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -6,7 +6,6 @@ exclude = (?x)(
           baybe/serialization
           | baybe/simulation
           | baybe/strategies
-          | baybe/surrogates
           | baybe/utils/dataframe.py
           | baybe/deprecation.py
           | baybe/exceptions.py

From 84c0bd4ecd15c565ec7e8589d7e5c47ee6750358 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 08:56:44 +0200
Subject: [PATCH 66/97] Ignore method overrides

---
 baybe/surrogates/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index 89b1dd24f..dcc39b792 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -73,8 +73,8 @@ def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
             _fit_original(self, train_x, train_y, context)
 
     # Replace the methods
-    cls._posterior = _posterior_new
-    cls._fit = _fit_new
+    cls._posterior = _posterior_new  # type: ignore
+    cls._fit = _fit_new  # type: ignore
 
     return cls
 

From 3a5909fc43728dbceacc0175733c564ab82faa32 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 08:57:10 +0200
Subject: [PATCH 67/97] Fix types of returned objects

---
 baybe/surrogates/gaussian_process/presets/default.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/presets/default.py b/baybe/surrogates/gaussian_process/presets/default.py
index a54c0705d..0e99e7dd4 100644
--- a/baybe/surrogates/gaussian_process/presets/default.py
+++ b/baybe/surrogates/gaussian_process/presets/default.py
@@ -96,16 +96,16 @@ def _default_noise_factory(
 
     # low D priors
     if train_x.shape[-1] < 10:  # <-- different condition compared to EDBO
-        return [GammaPrior(1.05, 0.5), 0.1]
+        return (GammaPrior(1.05, 0.5), 0.1)
 
     # DFT optimized priors
     elif uses_descriptors and train_x.shape[-1] < 100:
-        return [GammaPrior(1.5, 0.1), 5.0]
+        return (GammaPrior(1.5, 0.1), 5.0)
 
     # Mordred optimized priors
     elif uses_descriptors:
-        return [GammaPrior(1.5, 0.1), 5.0]
+        return (GammaPrior(1.5, 0.1), 5.0)
 
     # OHE optimized priors
     else:
-        return [GammaPrior(1.5, 0.1), 5.0]
+        return (GammaPrior(1.5, 0.1), 5.0)

From 8b4194b98b2fc669c28352f1bccd96c1f0022748 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 09:09:10 +0200
Subject: [PATCH 68/97] Add ngboost to mypy ignores

---
 mypy.ini | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mypy.ini b/mypy.ini
index 51a9b8bb3..e4946b313 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -25,6 +25,9 @@ ignore_missing_imports = True
 [mypy-mordred]
 ignore_missing_imports = True
 
+[mypy-ngboost]
+ignore_missing_imports = True
+
 [mypy-onnxruntime]
 ignore_missing_imports = True
 

From f9905f6c5f3bc0b0c53f9665952f5728a576eb98 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 09:09:30 +0200
Subject: [PATCH 69/97] Simplify sklearn mypy ignores

---
 mypy.ini | 17 +----------------
 1 file changed, 1 insertion(+), 16 deletions(-)

diff --git a/mypy.ini b/mypy.ini
index e4946b313..db2ede6e4 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -40,22 +40,7 @@ ignore_missing_imports = True
 [mypy-scipy.stats]
 ignore_missing_imports = True
 
-[mypy-sklearn.base]
-ignore_missing_imports = True
-
-[mypy-sklearn.cluster]
-ignore_missing_imports = True
-
-[mypy-sklearn.metrics]
-ignore_missing_imports = True
-
-[mypy-sklearn.mixture]
-ignore_missing_imports = True
-
-[mypy-sklearn.preprocessing]
-ignore_missing_imports = True
-
-[mypy-sklearn_extra.cluster]
+[mypy-sklearn.*]
 ignore_missing_imports = True
 
 [mypy-rdkit]

From 887bcef7906adc704a50a72504ad83f7354a3a9d Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 09:49:49 +0200
Subject: [PATCH 70/97] Fix model context typing

---
 baybe/surrogates/base.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index d26777cfb..6f5809816 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -4,7 +4,7 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, ClassVar, Protocol
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar
 
 import pandas as pd
 from attrs import define, field
@@ -48,6 +48,9 @@
 0-255 and the character set.
 """
 
+_ModelContext = TypeVar("_ModelContext", Any, None)
+"""Context information that is provided to the model fitting method."""
+
 
 class _NoTransform(Enum):
     """Sentinel class."""
@@ -85,7 +88,7 @@ def to_botorch(self) -> Model:
 
 
 @define
-class Surrogate(ABC, SurrogateProtocol, SerialMixin):
+class Surrogate(ABC, SurrogateProtocol, SerialMixin, Generic[_ModelContext]):
     """Abstract base class for all surrogate models."""
 
     # Class variables
@@ -248,7 +251,9 @@ def _posterior(self, candidates: Tensor, /) -> Posterior:
         """
 
     @staticmethod
-    def _get_model_context(searchspace: SearchSpace, objective: Objective) -> Any:
+    def _get_model_context(
+        searchspace: SearchSpace, objective: Objective
+    ) -> _ModelContext:
         """Get the surrogate-specific context for model fitting.
 
         By default, no context is created. If context is required, subclasses are
@@ -312,7 +317,7 @@ def fit(
         self._fit(train_x, train_y, self._get_model_context(searchspace, objective))
 
     @abstractmethod
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any = None) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None:
         """Perform the actual fitting logic."""
 
 

From cbf2f5aba5e62cf0c5ef5ef0c52e850b103377bc Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 09:51:12 +0200
Subject: [PATCH 71/97] Fix signature of CustomONNXSurrogate._fit

---
 baybe/surrogates/custom.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py
index 248bf73a6..ec7b55e74 100644
--- a/baybe/surrogates/custom.py
+++ b/baybe/surrogates/custom.py
@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar, NoReturn
+from typing import TYPE_CHECKING, Any, ClassVar, NoReturn
 
 from attrs import define, field, validators
 
@@ -96,7 +96,7 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
             torch.from_numpy(results[1]).pow(2).to(DTypeFloatTorch),
         )
 
-    def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         # TODO: This method actually needs to raise a NotImplementedError because
         #   ONNX surrogate models cannot be retrained. However, this would currently
         #   break the code since `BayesianRecommender` assumes that surrogates

From c13fa55e6f02aec1e90ae65530be0b58a420e84a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 10:23:51 +0200
Subject: [PATCH 72/97] Add explicit return values

---
 baybe/surrogates/gaussian_process/core.py | 4 ++--
 baybe/surrogates/ngboost.py               | 2 +-
 baybe/surrogates/random_forest.py         | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 5e9b22350..24a537126 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -131,14 +131,14 @@ def _make_parameter_scaler_factory(
         # See base class.
 
         # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
-        return
+        return None
 
     @staticmethod
     def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
         # See base class.
 
         # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
-        return
+        return None
 
     @staticmethod
     def _get_model_context(
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index b157ce65c..6da831878 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -60,7 +60,7 @@ def _make_parameter_scaler_factory(
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return
+        return None
 
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 07466c03f..06e868766 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -55,7 +55,7 @@ def _make_parameter_scaler_factory(
         # See base class.
 
         # Tree-like models do not require any input scaling
-        return
+        return None
 
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:

From 74b7e58d67b4c42ab5186645192f51d37abbe6cf Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 10:32:27 +0200
Subject: [PATCH 73/97] Disable output scaling for tree-based surrogates

---
 baybe/surrogates/ngboost.py       | 8 ++++++++
 baybe/surrogates/random_forest.py | 8 ++++++++
 2 files changed, 16 insertions(+)

diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 6da831878..8b53729a4 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -21,6 +21,7 @@
 
 if TYPE_CHECKING:
     from botorch.models.transforms.input import InputTransform
+    from botorch.models.transforms.outcome import OutcomeTransform
     from torch import Tensor
 
 
@@ -62,6 +63,13 @@ def _make_parameter_scaler_factory(
         # Tree-like models do not require any input scaling
         return None
 
+    @staticmethod
+    def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
+        # See base class.
+
+        # Tree-like models do not require any output scaling
+        return None
+
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 06e868766..967414e1b 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -22,6 +22,7 @@
 
 if TYPE_CHECKING:
     from botorch.models.transforms.input import InputTransform
+    from botorch.models.transforms.outcome import OutcomeTransform
     from torch import Tensor
 
 
@@ -57,6 +58,13 @@ def _make_parameter_scaler_factory(
         # Tree-like models do not require any input scaling
         return None
 
+    @staticmethod
+    def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
+        # See base class.
+
+        # Tree-like models do not require any output scaling
+        return None
+
     @batchify
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.

From 944ac8eaf9418e9a6d96c01dbecb2766cd1b994c Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 10:33:34 +0200
Subject: [PATCH 74/97] Add sklearn_extra to mypy ignores

---
 mypy.ini | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mypy.ini b/mypy.ini
index db2ede6e4..6a754cd8a 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -43,6 +43,9 @@ ignore_missing_imports = True
 [mypy-sklearn.*]
 ignore_missing_imports = True
 
+[mypy-sklearn_extra.*]
+ignore_missing_imports = True
+
 [mypy-rdkit]
 ignore_missing_imports = True
 

From 1175cd9d501494b3a81dd11bd1e3d6f179502406 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 10:42:05 +0200
Subject: [PATCH 75/97] Fix GP creation from preset

---
 baybe/surrogates/gaussian_process/core.py         | 2 +-
 baybe/surrogates/gaussian_process/presets/core.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 24a537126..bb3513ce7 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -114,7 +114,7 @@ class GaussianProcessSurrogate(Surrogate):
     _model = field(init=False, default=None, eq=False)
     """The actual model."""
 
-    @classmethod
+    @staticmethod
     def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate:
         """Create a Gaussian process surrogate from one of the defined presets."""
         return make_gp_from_preset(preset)
diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py
index df3276e63..8b06b313d 100644
--- a/baybe/surrogates/gaussian_process/presets/core.py
+++ b/baybe/surrogates/gaussian_process/presets/core.py
@@ -18,6 +18,8 @@ class GaussianProcessPreset(Enum):
 
 def make_gp_from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate:
     """Create a :class:`GaussianProcessSurrogate` from a :class:`GaussianProcessPreset."""  # noqa: E501
+    from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate
+
     if preset is GaussianProcessPreset.BAYBE:
         return GaussianProcessSurrogate()
 

From caa72ae01cad0ec2847fe4a11c3b13c77e07b855 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 10:43:23 +0200
Subject: [PATCH 76/97] Update path to Objective class

---
 baybe/surrogates/gaussian_process/core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index bb3513ce7..2a2952bbb 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -7,7 +7,7 @@
 from attrs import define, field
 from attrs.validators import instance_of
 
-from baybe.objective import Objective
+from baybe.objectives.base import Objective
 from baybe.parameters.base import Parameter
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.base import Surrogate

From 7a0f9e6dc69a23ba46c54eec08beb24887cccf95 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 10:52:35 +0200
Subject: [PATCH 77/97] Fix return type

---
 baybe/surrogates/gaussian_process/core.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 2a2952bbb..a0ae7c149 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -67,9 +67,9 @@ def parameter_bounds(self) -> Tensor:
 
         return torch.from_numpy(self.searchspace.comp_rep_bounds.values)
 
-    def get_numerical_indices(self, n_inputs: int) -> list[int]:
+    def get_numerical_indices(self, n_inputs: int) -> tuple[int, ...]:
         """Get the indices of the regular numerical model inputs."""
-        return [i for i in range(n_inputs) if i != self.task_idx]
+        return tuple(i for i in range(n_inputs) if i != self.task_idx)
 
 
 @define

From 67ccfd1e2aebbfdd106740a4a9cd8f13698fe137 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 10:58:28 +0200
Subject: [PATCH 78/97] Temporarily suppress mypy errors for batchify

There is no point of fixing typing at the moment since the
`joint_posterior` flag needs to be refactored/removed anyway.
The problem with the latter is that using surrogates that
offer no joint posterior in combination with continuous spaces
(or dense discrete spaces) makes not much sense, since they
cannot express correlation across closely neighboring points,
resulting in recommendation batches with parameter
configurations that are spatially highly concentrated.
---
 baybe/surrogates/utils.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index dcc39b792..5da61a317 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import typing
 from collections.abc import Callable
 from functools import wraps
 from typing import TYPE_CHECKING, Any
@@ -79,6 +80,9 @@ def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
     return cls
 
 
+# FIXME[typing]: Typing should be reactivated once the `joint_posterior` attribute
+#   has been refactored/removed
+@typing.no_type_check
 def batchify(
     posterior: Callable[[Surrogate, Tensor], tuple[Tensor, Tensor]],
 ) -> Callable[[Surrogate, Tensor], tuple[Tensor, Tensor]]:

From eaefde6bc5ae4004826f1c5cccd3d1daaac1e41d Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 11:40:32 +0200
Subject: [PATCH 79/97] Raise error when attempting to access posterior before
 training

---
 baybe/surrogates/base.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 6f5809816..f57e406fc 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -17,6 +17,7 @@
     UnstructureHook,
 )
 
+from baybe.exceptions import ModelNotTrainedError
 from baybe.objectives.base import Objective
 from baybe.parameters.base import Parameter
 from baybe.searchspace import SearchSpace
@@ -186,6 +187,9 @@ def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
             candidates: A dataframe containing parameter configurations in
                 **experimental representation**.
 
+        Raises:
+            ModelNotTrainedError: When called before the model has been trained.
+
         Returns:
             A :class:`botorch.posteriors.Posterior` object representing the posterior
             distribution at the given candidate points, where the posterior is also
@@ -193,6 +197,10 @@ def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
             lie in the same domain as the modelled targets/objective on which the
             surrogate was trained via :meth:`baybe.surrogates.base.Surrogate.fit`.
         """
+        if self._searchspace is None:
+            raise ModelNotTrainedError(
+                "The surrogate must be trained before a posterior can be computed."
+            )
         return self._posterior_comp_rep(
             to_tensor(self._searchspace.transform(candidates))
         )

From 12309782716963e53c6b5231a594fc74316f503b Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 11:41:48 +0200
Subject: [PATCH 80/97] Add typing workaround for accessing optional attributes

---
 baybe/surrogates/linear.py        | 4 ++++
 baybe/surrogates/ngboost.py       | 4 ++++
 baybe/surrogates/random_forest.py | 4 ++++
 3 files changed, 12 insertions(+)

diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py
index bc2f47ad1..b7e95a522 100644
--- a/baybe/surrogates/linear.py
+++ b/baybe/surrogates/linear.py
@@ -49,6 +49,10 @@ class BayesianLinearSurrogate(GaussianSurrogate):
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
 
+        # FIXME[typing]: It seems there is currently no better way to inform the type
+        #   checker that the attribute is available at the time of the function call
+        assert self._model is not None
+
         import torch
 
         # Get predictions
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index 8b53729a4..c7dfcb3ef 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -74,6 +74,10 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
 
+        # FIXME[typing]: It seems there is currently no better way to inform the type
+        #   checker that the attribute is available at the time of the function call
+        assert self._model is not None
+
         import torch
 
         # Get predictions
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 967414e1b..aa0c89d57 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -69,6 +69,10 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
     def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         # See base class.
 
+        # FIXME[typing]: It seems there is currently no better way to inform the type
+        #   checker that the attribute is available at the time of the function call
+        assert self._model is not None
+
         import torch
 
         # Evaluate all trees

From 08c3c2b5b9634d0378086f7d64193973b7640350 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 11:43:27 +0200
Subject: [PATCH 81/97] Clean up mypy.ini

---
 mypy.ini | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/mypy.ini b/mypy.ini
index 6a754cd8a..f89128faf 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -5,10 +5,7 @@ packages = baybe
 exclude = (?x)(
           baybe/serialization
           | baybe/simulation
-          | baybe/strategies
           | baybe/utils/dataframe.py
-          | baybe/deprecation.py
-          | baybe/exceptions.py
     	  | baybe/recommenders/naive.py
           | baybe/simulation.py
           )

From b1fd6f01f5e46b97775aa56f9dcbf2926a9d4d10 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Thu, 8 Aug 2024 17:33:14 +0200
Subject: [PATCH 82/97] Fix surrogate docstrings

---
 baybe/surrogates/base.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index f57e406fc..60a265f7b 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -233,16 +233,17 @@ def _posterior(self, candidates: Tensor, /) -> Posterior:
         This method is supposed to be overridden by subclasses to implement their
         model-specific surrogate architecture. Internally, the method is called by the
         base class with a **scaled** tensor of candidates in **computational
-        representation**, where the scaling is configurable by the subclass via its
-        other methods. The base class also takes care of transforming the returned
-        posterior back to the original scale according to the defined scalers.
+        representation**, where the scaling is configurable by the subclass by
+        overriding the default scaler factory methods of the base. The base class also
+        takes care of transforming the returned posterior back to the original scale
+        according to the defined scalers.
 
         This means:
         -----------
         Subclasses implementing this method do not have to bother about
         pre-/postprocessing of the in-/output. Instead, they only need to implement the
         mathematical operation of computing the posterior for the given input according
-        to their model specifications and can implicitly that scaling is handled
+        to their model specifications and can implicitly assume that scaling is handled
         appropriately outside. In short: the returned posterior simply needs to be on
         the same scale as the given input.
 

From 76013ac1424f07b0071ade274a76eee401129727 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 9 Aug 2024 10:24:25 +0200
Subject: [PATCH 83/97] Update ignore list in conf.py

---
 docs/conf.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/docs/conf.py b/docs/conf.py
index 804655aac..36b882157 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -116,12 +116,11 @@
     (r"py:.*", "baybe.constraints.conditions.Condition.__init__"),
     (r"py:.*", "baybe.serialization.mixin.SerialMixin.__init__"),
     (r"DeprecationWarning:", ""),
-    # Ignore the generics in utils.basic
-    # Might be able to us a regex here, is done explicitly at the moment for full
-    # transparency.
+    # Ignore the generics/aliases
     (r"py:class", "baybe.utils.basic._C"),
     (r"py:class", "baybe.utils.basic._T"),
     (r"py:class", "baybe.utils.basic._U"),
+    (r"ref:obj", "baybe.surrogates.base._ModelContext"),
     # Ignore custom class properties
     (r"py:obj", "baybe.acquisition.acqfs.*.is_mc"),
 ]

From b186360d25f73f652f93755d2783dc74f48a6ad5 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 9 Aug 2024 10:40:01 +0200
Subject: [PATCH 84/97] Make _ModelContext public to avoid sphinx problems

---
 baybe/surrogates/base.py | 8 ++++----
 docs/conf.py             | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 60a265f7b..1c4234bc5 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -49,7 +49,7 @@
 0-255 and the character set.
 """
 
-_ModelContext = TypeVar("_ModelContext", Any, None)
+ModelContext = TypeVar("ModelContext", Any, None)
 """Context information that is provided to the model fitting method."""
 
 
@@ -89,7 +89,7 @@ def to_botorch(self) -> Model:
 
 
 @define
-class Surrogate(ABC, SurrogateProtocol, SerialMixin, Generic[_ModelContext]):
+class Surrogate(ABC, SurrogateProtocol, SerialMixin, Generic[ModelContext]):
     """Abstract base class for all surrogate models."""
 
     # Class variables
@@ -262,7 +262,7 @@ def _posterior(self, candidates: Tensor, /) -> Posterior:
     @staticmethod
     def _get_model_context(
         searchspace: SearchSpace, objective: Objective
-    ) -> _ModelContext:
+    ) -> ModelContext:
         """Get the surrogate-specific context for model fitting.
 
         By default, no context is created. If context is required, subclasses are
@@ -326,7 +326,7 @@ def fit(
         self._fit(train_x, train_y, self._get_model_context(searchspace, objective))
 
     @abstractmethod
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor, context: ModelContext) -> None:
         """Perform the actual fitting logic."""
 
 
diff --git a/docs/conf.py b/docs/conf.py
index 36b882157..15c546c2e 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -120,7 +120,7 @@
     (r"py:class", "baybe.utils.basic._C"),
     (r"py:class", "baybe.utils.basic._T"),
     (r"py:class", "baybe.utils.basic._U"),
-    (r"ref:obj", "baybe.surrogates.base._ModelContext"),
+    (r"ref:obj", "baybe.surrogates.base.ModelContext"),
     # Ignore custom class properties
     (r"py:obj", "baybe.acquisition.acqfs.*.is_mc"),
 ]

From c00dee5d3390ffe32426653810df819ff7eaaad3 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 9 Aug 2024 11:44:55 +0200
Subject: [PATCH 85/97] Fix mypy error

---
 baybe/surrogates/gaussian_process/presets/default.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baybe/surrogates/gaussian_process/presets/default.py b/baybe/surrogates/gaussian_process/presets/default.py
index d84c6c5ba..bc5501dc8 100644
--- a/baybe/surrogates/gaussian_process/presets/default.py
+++ b/baybe/surrogates/gaussian_process/presets/default.py
@@ -87,5 +87,5 @@ def _default_noise_factory(
             np.interp(effective_dims, _DIM_LIMITS, [1.05, 1.5]),
             np.interp(effective_dims, _DIM_LIMITS, [0.5, 0.1]),
         ),
-        np.interp(effective_dims, _DIM_LIMITS, [0.1, 5.0]),
+        np.interp(effective_dims, _DIM_LIMITS, [0.1, 5.0]).item(),
     )

From c06cd46a74291d0476918e16bcd58e40ec6e051a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 9 Aug 2024 12:22:04 +0200
Subject: [PATCH 86/97] Allow extra columns in public posterior call

---
 baybe/surrogates/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 1c4234bc5..60801ce76 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -202,7 +202,7 @@ def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
                 "The surrogate must be trained before a posterior can be computed."
             )
         return self._posterior_comp_rep(
-            to_tensor(self._searchspace.transform(candidates))
+            to_tensor(self._searchspace.transform(candidates, allow_extra=True))
         )
 
     def _posterior_comp_rep(self, candidates: Tensor, /) -> Posterior:

From 1db3a9c0ccb1dba1a5ce934a1e1cc9ab6710fab9 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 9 Aug 2024 13:18:58 +0200
Subject: [PATCH 87/97] Add missing blank line to example

---
 examples/Custom_Hooks/campaign_stopping.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/Custom_Hooks/campaign_stopping.py b/examples/Custom_Hooks/campaign_stopping.py
index b5267b391..8b932baef 100644
--- a/examples/Custom_Hooks/campaign_stopping.py
+++ b/examples/Custom_Hooks/campaign_stopping.py
@@ -10,6 +10,7 @@
 # optimization if the found results are sufficiently good.
 # The underlying use case is taken from the example shown
 # [here](/examples/Backtesting/full_lookup).
+
 ### Imports
 import math
 import os

From a7b3b4e1501fdaa04d6676c05e753db91dda27c2 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 9 Aug 2024 13:35:46 +0200
Subject: [PATCH 88/97] Update use of surrogate in examples

---
 examples/Custom_Hooks/campaign_stopping.py          | 7 +++----
 examples/Custom_Hooks/probability_of_improvement.py | 6 +++---
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/examples/Custom_Hooks/campaign_stopping.py b/examples/Custom_Hooks/campaign_stopping.py
index 8b932baef..6adb9f07a 100644
--- a/examples/Custom_Hooks/campaign_stopping.py
+++ b/examples/Custom_Hooks/campaign_stopping.py
@@ -163,11 +163,10 @@ def stop_on_PI(
             f"Currently, only search spaces of type '{SearchSpaceType.DISCRETE}' are "
             f"accepted."
         )
-    train_x = searchspace.transform(measurements, allow_extra=True)
-    train_y = objective.transform(measurements)
     acqf = ProbabilityOfImprovement()
-
-    botorch_acqf = acqf.to_botorch(self.surrogate_model, searchspace, train_x, train_y)
+    botorch_acqf = acqf.to_botorch(
+        self.surrogate_model, searchspace, objective, measurements
+    )
     _, candidates_comp_rep = searchspace.discrete.get_candidates(
         allow_repeated_recommendations=self.allow_repeated_recommendations,
         allow_recommending_already_measured=self.allow_recommending_already_measured,
diff --git a/examples/Custom_Hooks/probability_of_improvement.py b/examples/Custom_Hooks/probability_of_improvement.py
index b53da8be2..f4e9f9afd 100644
--- a/examples/Custom_Hooks/probability_of_improvement.py
+++ b/examples/Custom_Hooks/probability_of_improvement.py
@@ -82,10 +82,10 @@ def extract_pi(
             f"Currently, only search spaces of type '{SearchSpaceType.DISCRETE}' are "
             f"accepted."
         )
-    train_x = searchspace.transform(measurements, allow_extra=True)
-    train_y = objective.transform(measurements)
     acqf = ProbabilityOfImprovement()
-    botorch_acqf = acqf.to_botorch(self.surrogate_model, searchspace, train_x, train_y)
+    botorch_acqf = acqf.to_botorch(
+        self.surrogate_model, searchspace, objective, measurements
+    )
     comp_rep_tensor = to_tensor(searchspace.discrete.comp_rep).unsqueeze(1)
     with torch.no_grad():
         pi = botorch_acqf(comp_rep_tensor)

From a1d37157f56378764109698a5c7207c72445f302 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Fri, 9 Aug 2024 13:52:05 +0200
Subject: [PATCH 89/97] Update CHANGELOG.md

---
 CHANGELOG.md | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 37a10e7f2..36c969762 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,8 +6,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Breaking Changes
-- `Surrogate` models now operate on dataframes in experimental representation instead of
-  tensors in computational representation
+- The public methods of `Surrogate` models now operate on dataframes in experimental
+  representation instead of tensors in computational representation
 - `Surrogate.posterior` models now returns a `Posterior` object
 - `param_bounds_comp` of `SearchSpace`, `SubspaceDiscrete` and `SubspaceContinuous` has
   been replaced with `comp_rep_bounds`, which returns a dataframe
@@ -16,15 +16,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `GaussianSurrogate` base class for surrogate models with Gaussian posteriors
 - `comp_rep_columns` property for `Parameter`, `SearchSpace`, `SubspaceDiscrete`
   and `SubspaceContinuous` classes
-- Reworked mechanisms for surrogate input/output scaling configurable per class
+- New mechanisms for surrogate input/output scaling configurable per class
 - `SurrogateProtocol` as an interface for user-defined surrogate architectures
 
 ### Changed
+- The transition from experimental to computational representation no longer happens
+  in the recommender but in the surrogate
 - Context information required by `Surrogate` models is now cleanly encapsulated into
   a `context` object passed to `Surrogate._fit`
-- Fallback models created by `catch_constant_targets` are stored outside of surrogate
+- Fallback models created by `catch_constant_targets` are stored outside the surrogate
 - `to_tensor` now also handles `numpy` arrays
-- `GaussianProcessSurrogate` no longer uses a separate scaling approach
 
 ### Removed
 - `register_custom_architecture` decorator

From d122d690e259f88ce61dd332b43cde674e627a79 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Tue, 13 Aug 2024 14:40:12 +0200
Subject: [PATCH 90/97] Indicate in docstring that scalers are fitted

---
 baybe/surrogates/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 60801ce76..079cac9cf 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -144,7 +144,7 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
         return Standardize
 
     def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
-        """Make the input scaler for transforming computational dataframes."""
+        """Make and fit the input scaler for transforming computational dataframes."""
         # Create a composite scaler from parameter-wise scaler objects
         mapping: dict[tuple[int, ...], InputTransform] = {}
         for p in searchspace.parameters:
@@ -163,7 +163,7 @@ def _make_input_scaler(self, searchspace: SearchSpace) -> ColumnTransformer:
     def _make_output_scaler(
         self, objective: Objective, measurements: pd.DataFrame
     ) -> OutcomeTransform | _NoTransform:
-        """Make the output scaler for transforming computational dataframes."""
+        """Make and fit the output scaler for transforming computational dataframes."""
         if (factory := self._make_target_scaler_factory()) is None:
             return _IDENTITY_TRANSFORM
 

From 1ad8d8f00f50e2524ad09ca46c21d1fe04e8b4f8 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 14 Aug 2024 12:34:38 +0200
Subject: [PATCH 91/97] Add backticks to docstring reference

---
 baybe/utils/scaling.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baybe/utils/scaling.py b/baybe/utils/scaling.py
index 4059e2402..0f4d5f2e9 100644
--- a/baybe/utils/scaling.py
+++ b/baybe/utils/scaling.py
@@ -25,7 +25,7 @@ class ColumnTransformer:
 
     @mapping.validator
     def _validate_mapping_types_lazily(self, attr, value):
-        """Perform transform isinstance check using lazy import."""
+        """Perform transform ``isinstance`` check using lazy import."""
         from botorch.models.transforms.input import InputTransform
 
         validator = deep_mapping(

From 4b22df6dbd8c67025e5c89a5429c9c330fa752dd Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 21 Aug 2024 08:44:23 +0200
Subject: [PATCH 92/97] Remove context argument from _fit signature

Since the searchspace needs to be stored during fitting anyway (because it is needed by the posterior method), we can simply use a regular attribute access and do not need to pass it via the context method argument
---
 CHANGELOG.md                              |  2 --
 baybe/surrogates/base.py                  | 22 ++++------------------
 baybe/surrogates/custom.py                |  4 ++--
 baybe/surrogates/gaussian_process/core.py | 12 +++---------
 baybe/surrogates/linear.py                |  2 +-
 baybe/surrogates/naive.py                 |  4 ++--
 baybe/surrogates/ngboost.py               |  2 +-
 baybe/surrogates/random_forest.py         |  2 +-
 baybe/surrogates/utils.py                 |  8 ++++----
 9 files changed, 18 insertions(+), 40 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 36c969762..43a6ed6ec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,8 +22,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 - The transition from experimental to computational representation no longer happens
   in the recommender but in the surrogate
-- Context information required by `Surrogate` models is now cleanly encapsulated into
-  a `context` object passed to `Surrogate._fit`
 - Fallback models created by `catch_constant_targets` are stored outside the surrogate
 - `to_tensor` now also handles `numpy` arrays
 
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 079cac9cf..2e8055202 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -4,7 +4,7 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar
+from typing import TYPE_CHECKING, ClassVar, Protocol
 
 import pandas as pd
 from attrs import define, field
@@ -49,9 +49,6 @@
 0-255 and the character set.
 """
 
-ModelContext = TypeVar("ModelContext", Any, None)
-"""Context information that is provided to the model fitting method."""
-
 
 class _NoTransform(Enum):
     """Sentinel class."""
@@ -89,7 +86,7 @@ def to_botorch(self) -> Model:
 
 
 @define
-class Surrogate(ABC, SurrogateProtocol, SerialMixin, Generic[ModelContext]):
+class Surrogate(ABC, SurrogateProtocol, SerialMixin):
     """Abstract base class for all surrogate models."""
 
     # Class variables
@@ -259,17 +256,6 @@ def _posterior(self, candidates: Tensor, /) -> Posterior:
             obtained via :meth:`baybe.surrogates.base.Surrogate._make_output_scaler`.
         """
 
-    @staticmethod
-    def _get_model_context(
-        searchspace: SearchSpace, objective: Objective
-    ) -> ModelContext:
-        """Get the surrogate-specific context for model fitting.
-
-        By default, no context is created. If context is required, subclasses are
-        expected to override this method.
-        """
-        return None
-
     def fit(
         self,
         searchspace: SearchSpace,
@@ -323,10 +309,10 @@ def fit(
             if self._output_scaler is _IDENTITY_TRANSFORM
             else self._output_scaler(train_y_comp_rep)[0]
         )
-        self._fit(train_x, train_y, self._get_model_context(searchspace, objective))
+        self._fit(train_x, train_y)
 
     @abstractmethod
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: ModelContext) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         """Perform the actual fitting logic."""
 
 
diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py
index ec7b55e74..eff3ecaa7 100644
--- a/baybe/surrogates/custom.py
+++ b/baybe/surrogates/custom.py
@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, NoReturn
+from typing import TYPE_CHECKING, ClassVar, NoReturn
 
 from attrs import define, field, validators
 
@@ -96,7 +96,7 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
             torch.from_numpy(results[1]).pow(2).to(DTypeFloatTorch),
         )
 
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         # TODO: This method actually needs to raise a NotImplementedError because
         #   ONNX surrogate models cannot be retrained. However, this would currently
         #   break the code since `BayesianRecommender` assumes that surrogates
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 307556646..89a6cf3da 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -7,7 +7,6 @@
 from attrs import define, field
 from attrs.validators import instance_of
 
-from baybe.objectives.base import Objective
 from baybe.parameters.base import Parameter
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.base import Surrogate
@@ -140,24 +139,19 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
         # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
         return None
 
-    @staticmethod
-    def _get_model_context(
-        searchspace: SearchSpace, objective: Objective
-    ) -> _ModelContext:
-        # See base class.
-        return _ModelContext(searchspace=searchspace)
-
     def _posterior(self, candidates: Tensor, /) -> Posterior:
         # See base class.
         return self._model.posterior(candidates)
 
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: _ModelContext) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         # See base class.
 
         import botorch
         import gpytorch
         import torch
 
+        context = _ModelContext(self._searchspace)
+
         numerical_idxs = context.get_numerical_indices(train_x.shape[-1])
 
         # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py
index b7e95a522..e979550a1 100644
--- a/baybe/surrogates/linear.py
+++ b/baybe/surrogates/linear.py
@@ -64,7 +64,7 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
 
         return mean, var
 
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         # See base class.
         self._model = ARDRegression(**(self.model_params))
         self._model.fit(train_x, train_y.ravel())
diff --git a/baybe/surrogates/naive.py b/baybe/surrogates/naive.py
index 201275321..5395ffb86 100644
--- a/baybe/surrogates/naive.py
+++ b/baybe/surrogates/naive.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 from attr import define, field
 
@@ -43,6 +43,6 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         var = torch.ones(len(candidates))
         return mean, var
 
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         # See base class.
         self._model = train_y.mean().item()
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index c7dfcb3ef..bc9db9666 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -89,6 +89,6 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
 
         return mean, var
 
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         # See base class.
         self._model = NGBRegressor(**(self.model_params)).fit(train_x, train_y.ravel())
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index aa0c89d57..00ece070d 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -94,7 +94,7 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
 
         return mean, var
 
-    def _fit(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+    def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         # See base class.
         self._model = RandomForestRegressor(**(self.model_params))
         self._model.fit(train_x, train_y.ravel())
diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index 5da61a317..a3d3d53ab 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -5,7 +5,7 @@
 import typing
 from collections.abc import Callable
 from functools import wraps
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from botorch.posteriors import Posterior
@@ -55,7 +55,7 @@ def _posterior_new(self, candidates: Tensor, /) -> Posterior:
         # Regular operation
         return _posterior_original(self, candidates)
 
-    def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
+    def _fit_new(self, train_x: Tensor, train_y: Tensor) -> None:
         """Original fit but with fallback model creation for constant targets."""
         if not (train_y.ndim == 2 and train_y.shape[-1] == 1):
             raise NotImplementedError(
@@ -65,13 +65,13 @@ def _fit_new(self, train_x: Tensor, train_y: Tensor, context: Any) -> None:
         # Alternative model fallback
         if train_y.numel() == 1 or train_y.std() < std_threshold:
             model = MeanPredictionSurrogate()
-            model._fit(train_x, train_y, context)
+            model._fit(train_x, train_y)
             _constant_target_model_store[id(self)] = model
 
         # Regular operation
         else:
             _constant_target_model_store.pop(id(self), None)
-            _fit_original(self, train_x, train_y, context)
+            _fit_original(self, train_x, train_y)
 
     # Replace the methods
     cls._posterior = _posterior_new  # type: ignore

From b74fbdfad736bcfb8fb013009c1160bc66ff016f Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Mon, 26 Aug 2024 10:34:59 +0200
Subject: [PATCH 93/97] Silence mypy error

---
 baybe/surrogates/gaussian_process/core.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 89a6cf3da..85fb082cd 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -150,6 +150,10 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         import gpytorch
         import torch
 
+        # FIXME[typing]: It seems there is currently no better way to inform the type
+        #   checker that the attribute is available at the time of the function call
+        assert self._searchspace is not None
+
         context = _ModelContext(self._searchspace)
 
         numerical_idxs = context.get_numerical_indices(train_x.shape[-1])

From 4f9e3cc6f01ed001bbfe288665db57aa624e97b0 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 28 Aug 2024 15:58:49 +0200
Subject: [PATCH 94/97] Replace generator comprehension with tuple in to_tensor

---
 baybe/utils/dataframe.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/baybe/utils/dataframe.py b/baybe/utils/dataframe.py
index 21737e281..a45ee8c0a 100644
--- a/baybe/utils/dataframe.py
+++ b/baybe/utils/dataframe.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Collection, Iterable, Iterator, Sequence
+from collections.abc import Collection, Iterable, Sequence
 from typing import (
     TYPE_CHECKING,
     Literal,
@@ -31,11 +31,11 @@ def to_tensor(x: np.ndarray | pd.DataFrame, /) -> Tensor: ...
 
 
 @overload
-def to_tensor(*x: np.ndarray | pd.DataFrame) -> Iterator[Tensor]: ...
+def to_tensor(*x: np.ndarray | pd.DataFrame) -> tuple[Tensor, ...]: ...
 
 
-def to_tensor(*x: np.ndarray | pd.DataFrame) -> Tensor | Iterator[Tensor]:
-    """Convert numpy arrays and pandas dataframes into tensors.
+def to_tensor(*x: np.ndarray | pd.DataFrame) -> Tensor | tuple[Tensor, ...]:
+    """Convert numpy arrays and pandas dataframes to tensors.
 
     Args:
         *x: The array(s)/dataframe(s) to be converted.
@@ -54,14 +54,14 @@ def to_tensor(*x: np.ndarray | pd.DataFrame) -> Tensor | Iterator[Tensor]:
 
     from baybe.utils.torch import DTypeFloatTorch
 
-    out = (
+    out = tuple(
         torch.from_numpy(
             (xi.values if isinstance(xi, pd.DataFrame) else xi).astype(DTypeFloatNumpy)
         ).to(DTypeFloatTorch)
         for xi in x
     )
     if len(x) == 1:
-        out = next(out)
+        out = out[0]
     return out
 
 

From 3b1635d962fa03926e4385d0d51c0be5a951f16a Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 28 Aug 2024 21:24:19 +0200
Subject: [PATCH 95/97] Indicate candidates domain using suffixes

---
 baybe/surrogates/base.py                  | 25 +++++++++++++----------
 baybe/surrogates/custom.py                |  6 ++++--
 baybe/surrogates/gaussian_process/core.py |  4 ++--
 baybe/surrogates/linear.py                |  6 ++++--
 baybe/surrogates/naive.py                 |  8 +++++---
 baybe/surrogates/ngboost.py               |  6 ++++--
 baybe/surrogates/random_forest.py         |  6 ++++--
 baybe/surrogates/utils.py                 |  6 +++---
 8 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index 2e8055202..c15611217 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -202,7 +202,7 @@ def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
             to_tensor(self._searchspace.transform(candidates, allow_extra=True))
         )
 
-    def _posterior_comp_rep(self, candidates: Tensor, /) -> Posterior:
+    def _posterior_comp_rep(self, candidates_comp: Tensor, /) -> Posterior:
         """Compute the posterior for candidates in computational representation.
 
         Takes a tensor of parameter configurations in **computational representation**
@@ -211,20 +211,20 @@ def _posterior_comp_rep(self, candidates: Tensor, /) -> Posterior:
         BoTorch's `optimize_*` functions.
 
         Args:
-            candidates: A tensor containing parameter configurations in **computational
-                representation**.
+            candidates_comp: A tensor containing parameter configurations in
+                **computational representation**.
 
         Returns:
             The same :class:`botorch.posteriors.Posterior` object as returned via
             :meth:`baybe.surrogates.base.Surrogate.posterior`.
         """
-        p = self._posterior(self._input_scaler.transform(candidates))
+        p = self._posterior(self._input_scaler.transform(candidates_comp))
         if self._output_scaler is not _IDENTITY_TRANSFORM:
             p = self._output_scaler.untransform_posterior(p)
         return p
 
     @abstractmethod
-    def _posterior(self, candidates: Tensor, /) -> Posterior:
+    def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior:
         """Perform the actual model-specific posterior evaluation logic.
 
         This method is supposed to be overridden by subclasses to implement their
@@ -245,9 +245,10 @@ def _posterior(self, candidates: Tensor, /) -> Posterior:
         the same scale as the given input.
 
         Args:
-            candidates: A tensor containing **pre-scaled** parameter configurations
-                in **computational representation**, as defined through the input scaler
-                obtained via :meth:`baybe.surrogates.base.Surrogate._make_input_scaler`.
+            candidates_comp_scaled: A tensor containing **pre-scaled** parameter
+                configurations in **computational representation**, as defined through
+                the input scaler obtained via
+                :meth:`baybe.surrogates.base.Surrogate._make_input_scaler`.
 
         Returns:
             A :class:`botorch.posteriors.Posterior` object representing the
@@ -320,7 +321,7 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
 class GaussianSurrogate(Surrogate, ABC):
     """A surrogate model providing Gaussian posterior estimates."""
 
-    def _posterior(self, candidates: Tensor, /) -> GPyTorchPosterior:
+    def _posterior(self, candidates_comp_scaled: Tensor, /) -> GPyTorchPosterior:
         # See base class.
 
         import torch
@@ -328,14 +329,16 @@ def _posterior(self, candidates: Tensor, /) -> GPyTorchPosterior:
         from gpytorch.distributions import MultivariateNormal
 
         # Construct the Gaussian posterior from the estimated first and second moment
-        mean, var = self._estimate_moments(candidates)
+        mean, var = self._estimate_moments(candidates_comp_scaled)
         if not self.joint_posterior:
             var = torch.diag_embed(var)
         mvn = MultivariateNormal(mean, var)
         return GPyTorchPosterior(mvn)
 
     @abstractmethod
-    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(
+        self, candidates_comp_scaled: Tensor, /
+    ) -> tuple[Tensor, Tensor]:
         """Estimate first and second moments of the Gaussian posterior.
 
         The second moment may either be a 1-D tensor of marginal variances for the
diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py
index eff3ecaa7..7bb0d425b 100644
--- a/baybe/surrogates/custom.py
+++ b/baybe/surrogates/custom.py
@@ -79,12 +79,14 @@ def default_model(self) -> ort.InferenceSession:
             raise ValueError("Invalid ONNX string") from exc
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(self, candidates_comp: Tensor, /) -> tuple[Tensor, Tensor]:
         import torch
 
         from baybe.utils.torch import DTypeFloatTorch
 
-        model_inputs = {self.onnx_input_name: candidates.numpy().astype(DTypeFloatONNX)}
+        model_inputs = {
+            self.onnx_input_name: candidates_comp.numpy().astype(DTypeFloatONNX)
+        }
         results = self._model.run(None, model_inputs)
 
         # IMPROVE: At the moment, we assume that the second model output contains
diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py
index 85fb082cd..ae8b244d2 100644
--- a/baybe/surrogates/gaussian_process/core.py
+++ b/baybe/surrogates/gaussian_process/core.py
@@ -139,9 +139,9 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
         # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above.
         return None
 
-    def _posterior(self, candidates: Tensor, /) -> Posterior:
+    def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior:
         # See base class.
-        return self._model.posterior(candidates)
+        return self._model.posterior(candidates_comp_scaled)
 
     def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         # See base class.
diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py
index e979550a1..f5272ed73 100644
--- a/baybe/surrogates/linear.py
+++ b/baybe/surrogates/linear.py
@@ -46,7 +46,9 @@ class BayesianLinearSurrogate(GaussianSurrogate):
     """The actual model."""
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(
+        self, candidates_comp_scaled: Tensor, /
+    ) -> tuple[Tensor, Tensor]:
         # See base class.
 
         # FIXME[typing]: It seems there is currently no better way to inform the type
@@ -56,7 +58,7 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         import torch
 
         # Get predictions
-        dists = self._model.predict(candidates.numpy(), return_std=True)
+        dists = self._model.predict(candidates_comp_scaled.numpy(), return_std=True)
 
         # Split into posterior mean and variance
         mean = torch.from_numpy(dists[0])
diff --git a/baybe/surrogates/naive.py b/baybe/surrogates/naive.py
index 5395ffb86..063d17549 100644
--- a/baybe/surrogates/naive.py
+++ b/baybe/surrogates/naive.py
@@ -33,14 +33,16 @@ class MeanPredictionSurrogate(GaussianSurrogate):
     """The estimated posterior mean value of the training targets."""
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(
+        self, candidates_comp_scaled: Tensor, /
+    ) -> tuple[Tensor, Tensor]:
         # See base class.
 
         import torch
 
         # TODO: use target value bounds for covariance scaling when explicitly provided
-        mean = self._model * torch.ones([len(candidates)])
-        var = torch.ones(len(candidates))
+        mean = self._model * torch.ones([len(candidates_comp_scaled)])
+        var = torch.ones(len(candidates_comp_scaled))
         return mean, var
 
     def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py
index bc9db9666..77296d587 100644
--- a/baybe/surrogates/ngboost.py
+++ b/baybe/surrogates/ngboost.py
@@ -71,7 +71,9 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
         return None
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(
+        self, candidates_comp_scaled: Tensor, /
+    ) -> tuple[Tensor, Tensor]:
         # See base class.
 
         # FIXME[typing]: It seems there is currently no better way to inform the type
@@ -81,7 +83,7 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         import torch
 
         # Get predictions
-        dists = self._model.pred_dist(candidates)
+        dists = self._model.pred_dist(candidates_comp_scaled)
 
         # Split into posterior mean and variance
         mean = torch.from_numpy(dists.mean())
diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py
index 00ece070d..888dc92a3 100644
--- a/baybe/surrogates/random_forest.py
+++ b/baybe/surrogates/random_forest.py
@@ -66,7 +66,9 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
         return None
 
     @batchify
-    def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
+    def _estimate_moments(
+        self, candidates_comp_scaled: Tensor, /
+    ) -> tuple[Tensor, Tensor]:
         # See base class.
 
         # FIXME[typing]: It seems there is currently no better way to inform the type
@@ -82,7 +84,7 @@ def _estimate_moments(self, candidates: Tensor, /) -> tuple[Tensor, Tensor]:
         predictions = torch.from_numpy(
             np.asarray(
                 [
-                    self._model.estimators_[tree].predict(candidates)
+                    self._model.estimators_[tree].predict(candidates_comp_scaled)
                     for tree in range(self._model.n_estimators)
                 ]
             )
diff --git a/baybe/surrogates/utils.py b/baybe/surrogates/utils.py
index a3d3d53ab..9f8b3c445 100644
--- a/baybe/surrogates/utils.py
+++ b/baybe/surrogates/utils.py
@@ -46,14 +46,14 @@ def catch_constant_targets(cls: type[Surrogate], std_threshold: float = 1e-6):
     _fit_original = cls._fit
     _posterior_original = cls._posterior
 
-    def _posterior_new(self, candidates: Tensor, /) -> Posterior:
+    def _posterior_new(self, candidates_comp_scaled: Tensor, /) -> Posterior:
         """Use fallback model if it exists, otherwise call original posterior."""
         # Alternative model fallback
         if constant_target_model := _constant_target_model_store.get(id(self), None):
-            return constant_target_model._posterior(candidates)
+            return constant_target_model._posterior(candidates_comp_scaled)
 
         # Regular operation
-        return _posterior_original(self, candidates)
+        return _posterior_original(self, candidates_comp_scaled)
 
     def _fit_new(self, train_x: Tensor, train_y: Tensor) -> None:
         """Original fit but with fallback model creation for constant targets."""

From 8e688c1ab9d67e6e60b06016e97252fa4f371888 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 28 Aug 2024 21:24:54 +0200
Subject: [PATCH 96/97] Rename _posterior_comp_rep to _posterior_comp

For consistency with its argument name
---
 baybe/surrogates/_adapter.py | 2 +-
 baybe/surrogates/base.py     | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/baybe/surrogates/_adapter.py b/baybe/surrogates/_adapter.py
index 6d167fe84..5ae2258b4 100644
--- a/baybe/surrogates/_adapter.py
+++ b/baybe/surrogates/_adapter.py
@@ -47,4 +47,4 @@ def posterior(  # noqa: D102
             raise NotImplementedError(
                 "The optional model posterior arguments are not yet implemented."
             )
-        return self._surrogate._posterior_comp_rep(X)
+        return self._surrogate._posterior_comp(X)
diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index c15611217..a0bc2bc67 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -198,11 +198,11 @@ def posterior(self, candidates: pd.DataFrame, /) -> Posterior:
             raise ModelNotTrainedError(
                 "The surrogate must be trained before a posterior can be computed."
             )
-        return self._posterior_comp_rep(
+        return self._posterior_comp(
             to_tensor(self._searchspace.transform(candidates, allow_extra=True))
         )
 
-    def _posterior_comp_rep(self, candidates_comp: Tensor, /) -> Posterior:
+    def _posterior_comp(self, candidates_comp: Tensor, /) -> Posterior:
         """Compute the posterior for candidates in computational representation.
 
         Takes a tensor of parameter configurations in **computational representation**

From ed32ab8f1524e1dea14572a30d8ba8f43b6e8423 Mon Sep 17 00:00:00 2001
From: AdrianSosic <adrian.sosic@merckgroup.com>
Date: Wed, 28 Aug 2024 21:27:35 +0200
Subject: [PATCH 97/97] Refine docstrings

---
 baybe/surrogates/base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py
index a0bc2bc67..b090e7920 100644
--- a/baybe/surrogates/base.py
+++ b/baybe/surrogates/base.py
@@ -123,7 +123,7 @@ def _make_parameter_scaler_factory(
         """Return the scaler factory to be used for the given parameter.
 
         This method is supposed to be overridden by subclasses to implement their
-        custom parameter scaling logic.
+        custom parameter scaling logic. Otherwise, parameters will be normalized.
         """
         from botorch.models.transforms.input import Normalize
 
@@ -134,7 +134,7 @@ def _make_target_scaler_factory() -> type[OutcomeTransform] | None:
         """Return the scaler factory to be used for target scaling.
 
         This method is supposed to be overridden by subclasses to implement their
-        custom target scaling logic.
+        custom target scaling logic. Otherwise, targets will be standardized.
         """
         from botorch.models.transforms.outcome import Standardize
 
@@ -245,7 +245,7 @@ def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior:
         the same scale as the given input.
 
         Args:
-            candidates_comp_scaled: A tensor containing **pre-scaled** parameter
+            candidates_comp_scaled: A tensor containing **scaled** parameter
                 configurations in **computational representation**, as defined through
                 the input scaler obtained via
                 :meth:`baybe.surrogates.base.Surrogate._make_input_scaler`.